Merge remote-tracking branch 'origin/dev' into fix-mcp-migrations

.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

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.
@@ -223,6 +224,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()
@@ -288,6 +291,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,

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

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

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py

@@ -52,7 +52,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.

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py

@@ -80,7 +80,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.

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py

@@ -80,7 +80,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.

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/mistral/adapter.py

@@ -69,7 +69,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.

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py

@@ -76,7 +76,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.
@@ -123,7 +123,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.
 
@@ -162,7 +162,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.
 

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py

@@ -112,7 +112,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.
@@ -154,6 +154,7 @@ class OpenAIAdapter(LLMInterface):
                     api_version=self.api_version,
                     response_model=response_model,
                     max_retries=self.MAX_RETRIES,
+                    **kwargs,
                 )
         except (
             ContentFilterFinishReasonError,
@@ -180,6 +181,7 @@ class OpenAIAdapter(LLMInterface):
                         # api_base=self.fallback_endpoint,
                         response_model=response_model,
                         max_retries=self.MAX_RETRIES,
+                        **kwargs,
                     )
             except (
                 ContentFilterFinishReasonError,
@@ -205,7 +207,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.
@@ -245,6 +247,7 @@ class OpenAIAdapter(LLMInterface):
             api_version=self.api_version,
             response_model=response_model,
             max_retries=self.MAX_RETRIES,
+            **kwargs,
         )
 
     @retry(
@@ -254,7 +257,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.
 
@@ -281,6 +284,7 @@ class OpenAIAdapter(LLMInterface):
                 api_base=self.endpoint,
                 api_version=self.api_version,
                 max_retries=self.MAX_RETRIES,
+                **kwargs,
             )
 
         return transcription
@@ -292,7 +296,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.
 
@@ -337,4 +341,5 @@ class OpenAIAdapter(LLMInterface):
             api_version=self.api_version,
             max_completion_tokens=300,
             max_retries=self.MAX_RETRIES,
+            **kwargs,
         )

cognee/tasks/graph/extract_graph_from_data.py

@@ -97,6 +97,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.
@@ -111,7 +112,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
         ]
     )