From 0a4b1068a253df8fb4e39a93ee18a73c911ee49e Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Mon, 17 Nov 2025 17:42:22 +0100 Subject: [PATCH 1/5] feat: add kwargs to openai adapter functions --- .../litellm_instructor/llm/openai/adapter.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py index 305b426b8..152f43e33 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py @@ -108,7 +108,7 @@ class OpenAIAdapter(LLMInterface): reraise=True, ) async def acreate_structured_output( - self, text_input: str, system_prompt: str, response_model: Type[BaseModel] + self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs ) -> BaseModel: """ Generate a response from a user query. @@ -149,6 +149,7 @@ class OpenAIAdapter(LLMInterface): api_version=self.api_version, response_model=response_model, max_retries=self.MAX_RETRIES, + **kwargs, ) except ( ContentFilterFinishReasonError, @@ -174,6 +175,7 @@ class OpenAIAdapter(LLMInterface): # api_base=self.fallback_endpoint, response_model=response_model, max_retries=self.MAX_RETRIES, + **kwargs, ) except ( ContentFilterFinishReasonError, @@ -199,7 +201,7 @@ class OpenAIAdapter(LLMInterface): reraise=True, ) def create_structured_output( - self, text_input: str, system_prompt: str, response_model: Type[BaseModel] + self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs ) -> BaseModel: """ Generate a response from a user query. @@ -239,6 +241,7 @@ class OpenAIAdapter(LLMInterface): api_version=self.api_version, response_model=response_model, max_retries=self.MAX_RETRIES, + **kwargs, ) @retry( @@ -248,7 +251,7 @@ class OpenAIAdapter(LLMInterface): before_sleep=before_sleep_log(logger, logging.DEBUG), reraise=True, ) - async def create_transcript(self, input): + async def create_transcript(self, input, **kwargs): """ Generate an audio transcript from a user query. @@ -275,6 +278,7 @@ class OpenAIAdapter(LLMInterface): api_base=self.endpoint, api_version=self.api_version, max_retries=self.MAX_RETRIES, + **kwargs, ) return transcription @@ -286,7 +290,7 @@ class OpenAIAdapter(LLMInterface): before_sleep=before_sleep_log(logger, logging.DEBUG), reraise=True, ) - async def transcribe_image(self, input) -> BaseModel: + async def transcribe_image(self, input, **kwargs) -> BaseModel: """ Generate a transcription of an image from a user query. @@ -331,4 +335,5 @@ class OpenAIAdapter(LLMInterface): api_version=self.api_version, max_completion_tokens=300, max_retries=self.MAX_RETRIES, + **kwargs, ) From aa8afefe8a7ae4233e82edc71ee9441f0b68d325 Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Thu, 27 Nov 2025 17:05:37 +0100 Subject: [PATCH 2/5] feat: add kwargs to cognify and related tasks --- cognee/api/v1/cognify/cognify.py | 4 ++++ cognee/infrastructure/llm/LLMGateway.py | 4 ++-- .../llm/extraction/knowledge_graph/extract_content_graph.py | 4 ++-- cognee/tasks/graph/extract_graph_from_data.py | 3 ++- 4 files changed, 10 insertions(+), 5 deletions(-) diff --git a/cognee/api/v1/cognify/cognify.py b/cognee/api/v1/cognify/cognify.py index 0fa345176..bb2ebe86e 100644 --- a/cognee/api/v1/cognify/cognify.py +++ b/cognee/api/v1/cognify/cognify.py @@ -53,6 +53,7 @@ async def cognify( custom_prompt: Optional[str] = None, temporal_cognify: bool = False, data_per_batch: int = 20, + **kwargs ): """ Transform ingested data into a structured knowledge graph. @@ -224,6 +225,7 @@ async def cognify( config=config, custom_prompt=custom_prompt, chunks_per_batch=chunks_per_batch, + **kwargs, ) # By calling get pipeline executor we get a function that will have the run_pipeline run in the background or a function that we will need to wait for @@ -251,6 +253,7 @@ async def get_default_tasks( # TODO: Find out a better way to do this (Boris's config: Config = None, custom_prompt: Optional[str] = None, chunks_per_batch: int = 100, + **kwargs, ) -> list[Task]: if config is None: ontology_config = get_ontology_env_config() @@ -286,6 +289,7 @@ async def get_default_tasks( # TODO: Find out a better way to do this (Boris's config=config, custom_prompt=custom_prompt, task_config={"batch_size": chunks_per_batch}, + **kwargs, ), # Generate knowledge graphs from the document chunks. Task( summarize_text, diff --git a/cognee/infrastructure/llm/LLMGateway.py b/cognee/infrastructure/llm/LLMGateway.py index ab5bb35d7..fd42eb55e 100644 --- a/cognee/infrastructure/llm/LLMGateway.py +++ b/cognee/infrastructure/llm/LLMGateway.py @@ -11,7 +11,7 @@ class LLMGateway: @staticmethod def acreate_structured_output( - text_input: str, system_prompt: str, response_model: Type[BaseModel] + text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs ) -> Coroutine: llm_config = get_llm_config() if llm_config.structured_output_framework.upper() == "BAML": @@ -31,7 +31,7 @@ class LLMGateway: llm_client = get_llm_client() return llm_client.acreate_structured_output( - text_input=text_input, system_prompt=system_prompt, response_model=response_model + text_input=text_input, system_prompt=system_prompt, response_model=response_model, **kwargs ) @staticmethod diff --git a/cognee/infrastructure/llm/extraction/knowledge_graph/extract_content_graph.py b/cognee/infrastructure/llm/extraction/knowledge_graph/extract_content_graph.py index 59e6f563a..4a40979f4 100644 --- a/cognee/infrastructure/llm/extraction/knowledge_graph/extract_content_graph.py +++ b/cognee/infrastructure/llm/extraction/knowledge_graph/extract_content_graph.py @@ -10,7 +10,7 @@ from cognee.infrastructure.llm.config import ( async def extract_content_graph( - content: str, response_model: Type[BaseModel], custom_prompt: Optional[str] = None + content: str, response_model: Type[BaseModel], custom_prompt: Optional[str] = None, **kwargs ): if custom_prompt: system_prompt = custom_prompt @@ -30,7 +30,7 @@ async def extract_content_graph( system_prompt = render_prompt(prompt_path, {}, base_directory=base_directory) content_graph = await LLMGateway.acreate_structured_output( - content, system_prompt, response_model + content, system_prompt, response_model, **kwargs ) return content_graph diff --git a/cognee/tasks/graph/extract_graph_from_data.py b/cognee/tasks/graph/extract_graph_from_data.py index 49b51af2d..965214677 100644 --- a/cognee/tasks/graph/extract_graph_from_data.py +++ b/cognee/tasks/graph/extract_graph_from_data.py @@ -99,6 +99,7 @@ async def extract_graph_from_data( graph_model: Type[BaseModel], config: Config = None, custom_prompt: Optional[str] = None, + **kwargs, ) -> List[DocumentChunk]: """ Extracts and integrates a knowledge graph from the text content of document chunks using a specified graph model. @@ -113,7 +114,7 @@ async def extract_graph_from_data( chunk_graphs = await asyncio.gather( *[ - extract_content_graph(chunk.text, graph_model, custom_prompt=custom_prompt) + extract_content_graph(chunk.text, graph_model, custom_prompt=custom_prompt, **kwargs) for chunk in data_chunks ] ) From af8c5bedcc48e18c3723a2fbfa8afba3de242cbb Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Thu, 11 Dec 2025 17:47:23 +0100 Subject: [PATCH 3/5] feat: add kwargs to other adapters --- .../litellm_instructor/llm/anthropic/adapter.py | 2 +- .../litellm_instructor/llm/gemini/adapter.py | 2 +- .../litellm_instructor/llm/generic_llm_api/adapter.py | 2 +- .../litellm_instructor/llm/mistral/adapter.py | 2 +- .../litellm_instructor/llm/ollama/adapter.py | 6 +++--- 5 files changed, 7 insertions(+), 7 deletions(-) diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py index dbf0dfbea..46e2b2736 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py @@ -51,7 +51,7 @@ class AnthropicAdapter(LLMInterface): reraise=True, ) async def acreate_structured_output( - self, text_input: str, system_prompt: str, response_model: Type[BaseModel] + self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs ) -> BaseModel: """ Generate a response from a user query. diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py index 226f291d7..66d53b842 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py @@ -79,7 +79,7 @@ class GeminiAdapter(LLMInterface): reraise=True, ) async def acreate_structured_output( - self, text_input: str, system_prompt: str, response_model: Type[BaseModel] + self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs ) -> BaseModel: """ Generate a response from a user query. diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py index 9d7f25fc5..3049b3c4f 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py @@ -79,7 +79,7 @@ class GenericAPIAdapter(LLMInterface): reraise=True, ) async def acreate_structured_output( - self, text_input: str, system_prompt: str, response_model: Type[BaseModel] + self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs ) -> BaseModel: """ Generate a response from a user query. diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/mistral/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/mistral/adapter.py index 355cdae0b..146d0a07a 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/mistral/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/mistral/adapter.py @@ -68,7 +68,7 @@ class MistralAdapter(LLMInterface): reraise=True, ) async def acreate_structured_output( - self, text_input: str, system_prompt: str, response_model: Type[BaseModel] + self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs ) -> BaseModel: """ Generate a response from the user query. diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py index aabd19867..5ae09a4ac 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py @@ -74,7 +74,7 @@ class OllamaAPIAdapter(LLMInterface): reraise=True, ) async def acreate_structured_output( - self, text_input: str, system_prompt: str, response_model: Type[BaseModel] + self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs ) -> BaseModel: """ Generate a structured output from the LLM using the provided text and system prompt. @@ -121,7 +121,7 @@ class OllamaAPIAdapter(LLMInterface): before_sleep=before_sleep_log(logger, logging.DEBUG), reraise=True, ) - async def create_transcript(self, input_file: str) -> str: + async def create_transcript(self, input_file: str, **kwargs) -> str: """ Generate an audio transcript from a user query. @@ -160,7 +160,7 @@ class OllamaAPIAdapter(LLMInterface): before_sleep=before_sleep_log(logger, logging.DEBUG), reraise=True, ) - async def transcribe_image(self, input_file: str) -> str: + async def transcribe_image(self, input_file: str, **kwargs) -> str: """ Transcribe content from an image using base64 encoding. From 14ff94f269599140df6e830761ef3b6f2c99eb28 Mon Sep 17 00:00:00 2001 From: Pavel Zorin Date: Thu, 11 Dec 2025 12:38:19 +0100 Subject: [PATCH 4/5] Initial release pipeline --- .github/workflows/release.yml | 154 ++++++++++++++++++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 .github/workflows/release.yml diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 000000000..a19635628 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,154 @@ +name: release.yml +on: + workflow_dispatch: + inputs: + flavour: + required: true + default: dev + type: choice + options: + - dev + - main + description: Dev or Main release + test_mode: + required: true + type: boolean + description: Aka Dry Run. If true, it won't affect public indices or repositories + +jobs: + release-github: + name: Create GitHub Release from ${{ inputs.flavour }} + outputs: + tag: ${{ steps.create_tag.outputs.tag }} + version: ${{ steps.create_tag.outputs.version }} + permissions: + contents: write + runs-on: ubuntu-latest + + steps: + - name: Check out ${{ inputs.flavour }} + uses: actions/checkout@v4 + with: + ref: ${{ inputs.flavour }} + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Create and push git tag + id: create_tag + env: + TEST_MODE: ${{ inputs.test_mode }} + run: | + VERSION="$(uv version --short)" + TAG="v${VERSION}" + + echo "Tag to create: ${TAG}" + + git config user.name "github-actions[bot]" + git config user.email "41898282+github-actions[bot]@users.noreply.github.com" + + echo "tag=${TAG}" >> "$GITHUB_OUTPUT" + echo "version=${VERSION}" >> "$GITHUB_OUTPUT" + + if [ "$TEST_MODE" = "false" ]; then + git tag "${TAG}" + git push origin "${TAG}" + else + echo "Test mode is enabled. Skipping tag creation and push." + fi + + - name: Create GitHub Release + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ steps.create_tag.outputs.tag }} + generate_release_notes: true + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + release-pypi-package: + needs: release-github + name: Release PyPI Package from ${{ inputs.flavour }} + permissions: + contents: read + runs-on: ubuntu-latest + + steps: + - name: Check out ${{ inputs.flavour }} + uses: actions/checkout@v4 + with: + ref: ${{ inputs.flavour }} + + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Install Python + run: uv python install + + - name: Install dependencies + run: uv sync --locked --all-extras + + - name: Build distributions + run: uv build + + - name: Publish ${{ inputs.flavour }} release to TestPyPI + if: ${{ !inputs.test_mode }} + env: + UV_PUBLISH_TOKEN: ${{ secrets.TEST_PYPI_TOKEN }} + run: uv publish --publish-url https://test.pypi.org/legacy/ + + - name: Publish ${{ inputs.flavour }} release to PyPI + if: ${{ !inputs.test_mode }} + env: + UV_PUBLISH_TOKEN: ${{ secrets.PYPI_TOKEN }} + run: uv publish + + release-docker-image: + needs: release-github + name: Release Docker Image from ${{ inputs.flavour }} + permissions: + contents: read + runs-on: ubuntu-latest + + steps: + - name: Check out ${{ inputs.flavour }} + uses: actions/checkout@v4 + with: + ref: ${{ inputs.flavour }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_PASSWORD }} + + - name: Build and push Dev Docker Image + if: ${{ inputs.flavour == 'dev' }} + uses: docker/build-push-action@v5 + with: + context: . + platforms: linux/amd64,linux/arm64 + push: ${{ !inputs.test_mode }} + tags: cognee/cognee:${{ needs.release-github.outputs.version }} + labels: | + version=${{ needs.release-github.outputs.version }} + flavour=${{ inputs.flavour }} + cache-from: type=registry,ref=cognee/cognee:buildcache + cache-to: type=registry,ref=cognee/cognee:buildcache,mode=max + + - name: Build and push Main Docker Image + if: ${{ inputs.flavour == 'main' }} + uses: docker/build-push-action@v5 + with: + context: . + platforms: linux/amd64,linux/arm64 + push: ${{ !inputs.test_mode }} + tags: | + cognee/cognee:${{ needs.release-github.outputs.version }} + cognee/cognee:latest + labels: | + version=${{ needs.release-github.outputs.version }} + flavour=${{ inputs.flavour }} + cache-from: type=registry,ref=cognee/cognee:buildcache + cache-to: type=registry,ref=cognee/cognee:buildcache,mode=max From a6bc27afaaeb901e5e771a84ca5e9ba2af473aba Mon Sep 17 00:00:00 2001 From: Pavel Zorin Date: Fri, 12 Dec 2025 17:31:54 +0100 Subject: [PATCH 5/5] Cleanup --- .github/workflows/release.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a19635628..ff2f809f3 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -90,7 +90,7 @@ jobs: run: uv build - name: Publish ${{ inputs.flavour }} release to TestPyPI - if: ${{ !inputs.test_mode }} + if: ${{ inputs.test_mode }} env: UV_PUBLISH_TOKEN: ${{ secrets.TEST_PYPI_TOKEN }} run: uv publish --publish-url https://test.pypi.org/legacy/