cognee 0.3.3__tar.gz → 0.3.4.dev0__tar.gz

{cognee-0.3.3 → cognee-0.3.4.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.3.3
+Version: 0.3.4.dev0
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee
@@ -57,7 +57,7 @@ Requires-Dist: structlog<26,>=25.2.0
 Requires-Dist: tiktoken<1.0.0,>=0.8.0
 Requires-Dist: typing-extensions<5.0.0,>=4.12.2
 Provides-Extra: anthropic
-Requires-Dist: anthropic<0.27,>=0.26.1; extra == 'anthropic'
+Requires-Dist: anthropic>=0.27; extra == 'anthropic'
 Provides-Extra: api
 Requires-Dist: gunicorn<24,>=20.1.0; extra == 'api'
 Requires-Dist: uvicorn<1.0.0,>=0.34.0; extra == 'api'
@@ -65,8 +65,8 @@ Requires-Dist: websockets<16.0.0,>=15.0.1; extra == 'api'
 Provides-Extra: aws
 Requires-Dist: s3fs[boto3]==2025.3.2; extra == 'aws'
 Provides-Extra: chromadb
-Requires-Dist: chromadb<0.7,>=0.3.0; extra == 'chromadb'
-Requires-Dist: pypika==0.48.8; extra == 'chromadb'
+Requires-Dist: chromadb<0.7,>=0.6; extra == 'chromadb'
+Requires-Dist: pypika==0.48.9; extra == 'chromadb'
 Provides-Extra: codegraph
 Requires-Dist: fastembed<=0.6.0; (python_version < '3.13') and extra == 'codegraph'
 Requires-Dist: transformers<5,>=4.46.3; extra == 'codegraph'
@@ -316,7 +316,7 @@ You can also cognify your files and query using cognee UI.
 
 <img src="assets/cognee-new-ui.webp" width="100%" alt="Cognee UI 2"></a>
 
-Try cognee UI by runnning ``` cognee -ui ``` command on your terminal.
+Try cognee UI by runnning ``` cognee-cli -ui ``` command on your terminal.
 
 ## Understand our architecture
 

{cognee-0.3.3 → cognee-0.3.4.dev0}/README.md

@@ -176,7 +176,7 @@ You can also cognify your files and query using cognee UI.
 
 <img src="assets/cognee-new-ui.webp" width="100%" alt="Cognee UI 2"></a>
 
-Try cognee UI by runnning ``` cognee -ui ``` command on your terminal.
+Try cognee UI by runnning ``` cognee-cli -ui ``` command on your terminal.
 
 ## Understand our architecture
 

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/api/v1/search/search.py

@@ -22,7 +22,7 @@ async def search(
     node_type: Optional[Type] = NodeSet,
     node_name: Optional[List[str]] = None,
     save_interaction: bool = False,
-    last_k: Optional[int] = None,
+    last_k: Optional[int] = 1,
     only_context: bool = False,
     use_combined_context: bool = False,
 ) -> Union[List[SearchResult], CombinedSearchResult]:

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py

@@ -83,7 +83,7 @@ def process_data_for_chroma(data):
         elif isinstance(value, list):
             # Store lists as JSON strings with special prefix
             processed_data[f"{key}__list"] = json.dumps(value)
-        elif isinstance(value, (str, int, float, bool)) or value is None:
+        elif isinstance(value, (str, int, float, bool)):
             processed_data[key] = value
         else:
             processed_data[key] = str(value)
@@ -553,8 +553,4 @@ class ChromaDBAdapter(VectorDBInterface):
             Returns a list of collection names.
         """
         client = await self.get_connection()
-        collections = await client.list_collections()
-        return [
-            collection.name if hasattr(collection, "name") else collection["name"]
-            for collection in collections
-        ]
+        return await client.list_collections()

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

@@ -94,10 +94,8 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         """
         Internal method to call the Ollama embeddings endpoint for a single prompt.
         """
-        payload = {
-            "model": self.model,
-            "prompt": prompt,
-        }
+        payload = {"model": self.model, "prompt": prompt, "input": prompt}
+
         headers = {}
         api_key = os.getenv("LLM_API_KEY")
         if api_key:

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py

@@ -12,6 +12,7 @@ from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.ll
 )
 
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.infrastructure.llm.config import get_llm_config
 
 
 class AnthropicAdapter(LLMInterface):
@@ -27,7 +28,8 @@ class AnthropicAdapter(LLMInterface):
         import anthropic
 
         self.aclient = instructor.patch(
-            create=anthropic.AsyncAnthropic().messages.create, mode=instructor.Mode.ANTHROPIC_TOOLS
+            create=anthropic.AsyncAnthropic(api_key=get_llm_config().llm_api_key).messages.create,
+            mode=instructor.Mode.ANTHROPIC_TOOLS,
         )
 
         self.model = model
@@ -57,7 +59,7 @@ class AnthropicAdapter(LLMInterface):
 
         return await self.aclient(
             model=self.model,
-            max_completion_tokens=4096,
+            max_tokens=4096,
             max_retries=5,
             messages=[
                 {

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/modules/notebooks/methods/create_tutorial_notebook.py

@@ -1,4 +1,3 @@
-
 from uuid import UUID, uuid4
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -28,14 +27,14 @@ async def create_tutorial_notebook(user_id: UUID, session: AsyncSession):
         session=session,
     )
 
+
 cell_content = [
-"""
+    """
 # Using Cognee with Python Development Data
 
 Unite authoritative Python practice (Guido van Rossum's own contributions!), normative guidance (Zen/PEP 8), and your lived context (rules + conversations) into one *AI memory* that produces answers that are relevant, explainable, and consistent.
 """,
-
-"""
+    """
 ## What You'll Learn
 
 In this comprehensive tutorial, you'll discover how to transform scattered development data into an intelligent knowledge system that enhances your coding workflow. By the end, you'll have:
@@ -46,8 +45,7 @@ In this comprehensive tutorial, you'll discover how to transform scattered devel
 
 This tutorial demonstrates the power of **knowledge graphs** and **retrieval-augmented generation (RAG)** for software development, showing you how to build systems that learn from Python's creator and improve your own Python development.
 """,
-
-"""
+    """
 ## Cognee and its core operations
 
 Before we dive in, let's understand the core Cognee operations we'll be working with:
@@ -56,8 +54,7 @@ Before we dive in, let's understand the core Cognee operations we'll be working
 - `cognee.search()` - Queries the knowledge graph with natural language or Cypher
 - `cognee.memify()` - Cognee's \"secret sauce\" that infers implicit connections and rules from your data
 """,
-
-"""
+    """
 ## Data used in this tutorial
 
 Cognee can ingest many types of sources. In this tutorial, we use a small, concrete set of files that cover different perspectives:
@@ -67,8 +64,7 @@ Cognee can ingest many types of sources. In this tutorial, we use a small, concr
 - `my_developer_rules.md` — Local constraints. Your house rules, conventions, and project-specific requirements (scope, privacy, Spec.md). Keeps recommendations relevant to your actual workflow.
 - `copilot_conversations.json` — Personal history. Transcripts of real assistant conversations, including your questions, code snippets, and discussion topics. Captures "how you code" and connects it to "how Guido codes."
 """,
-
-"""
+    """
 # Preliminaries
 
 To strike the balanace between speed, cost, anc quality, we recommend using OpenAI's `4o-mini` model; make sure your `.env` file contains this line:
@@ -76,8 +72,7 @@ To strike the balanace between speed, cost, anc quality, we recommend using Open
 LLM_MODEL="gpt-4o-mini"
 `
 """,
-
-"""
+    """
 import cognee
 
 result = await cognee.add(
@@ -88,5 +83,5 @@ result = await cognee.add(
 await cognee.cognify(temporal_cognify=True)
 
 results = await cognee.search("Show me commits")
-"""
-]
+""",
+]

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/modules/search/methods/search.py

@@ -136,12 +136,19 @@ async def search(
         if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true":
             return_value = []
             for search_result in search_results:
-                result, context, datasets = search_result
+                prepared_search_results = await prepare_search_result(search_result)
+
+                result = prepared_search_results["result"]
+                graphs = prepared_search_results["graphs"]
+                context = prepared_search_results["context"]
+                datasets = prepared_search_results["datasets"]
+
                 return_value.append(
                     {
-                        "search_result": result,
+                        "search_result": [result] if result else None,
                         "dataset_id": datasets[0].id,
                         "dataset_name": datasets[0].name,
+                        "graphs": graphs,
                     }
                 )
             return return_value
@@ -155,14 +162,6 @@ async def search(
                 return return_value[0]
             else:
                 return return_value
-        # return [
-        #     SearchResult(
-        #         search_result=result,
-        #         dataset_id=datasets[min(index, len(datasets) - 1)].id if datasets else None,
-        #         dataset_name=datasets[min(index, len(datasets) - 1)].name if datasets else None,
-        #     )
-        #     for index, (result, _, datasets) in enumerate(search_results)
-        # ]
 
 
 async def authorized_search(
@@ -208,11 +207,11 @@ async def authorized_search(
         context = {}
         datasets: List[Dataset] = []
 
-        for _, search_context, datasets in search_responses:
-            for dataset in datasets:
+        for _, search_context, search_datasets in search_responses:
+            for dataset in search_datasets:
                 context[str(dataset.id)] = search_context
 
-            datasets.extend(datasets)
+            datasets.extend(search_datasets)
 
         specific_search_tools = await get_search_type_tools(
             query_type=query_type,

cognee-0.3.4.dev0/cognee/modules/search/utils/prepare_search_result.py

@@ -0,0 +1,63 @@
+from typing import List, cast
+from uuid import uuid5, NAMESPACE_OID
+
+from cognee.modules.graph.utils import resolve_edges_to_text
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+from cognee.modules.search.types.SearchResult import SearchResultDataset
+from cognee.modules.search.utils.transform_context_to_graph import transform_context_to_graph
+from cognee.modules.search.utils.transform_insights_to_graph import transform_insights_to_graph
+
+
+async def prepare_search_result(search_result):
+    results, context, datasets = search_result
+
+    graphs = None
+    result_graph = None
+    context_texts = {}
+
+    if isinstance(datasets, list) and len(datasets) == 0:
+        datasets = [
+            SearchResultDataset(
+                id=uuid5(NAMESPACE_OID, "*"),
+                name="all available datasets",
+            )
+        ]
+
+    if (
+        isinstance(context, List)
+        and len(context) > 0
+        and isinstance(context[0], tuple)
+        and context[0][1].get("relationship_name")
+    ):
+        context_graph = transform_insights_to_graph(context)
+        graphs = {
+            ", ".join([dataset.name for dataset in datasets]): context_graph,
+        }
+        results = None
+    elif isinstance(context, List) and len(context) > 0 and isinstance(context[0], Edge):
+        context_graph = transform_context_to_graph(context)
+
+        graphs = {
+            ", ".join([dataset.name for dataset in datasets]): context_graph,
+        }
+        context_texts = {
+            ", ".join([dataset.name for dataset in datasets]): await resolve_edges_to_text(context),
+        }
+    elif isinstance(context, str):
+        context_texts = {
+            ", ".join([dataset.name for dataset in datasets]): context,
+        }
+    elif isinstance(context, List) and len(context) > 0 and isinstance(context[0], str):
+        context_texts = {
+            ", ".join([dataset.name for dataset in datasets]): "\n".join(cast(List[str], context)),
+        }
+
+    if isinstance(results, List) and len(results) > 0 and isinstance(results[0], Edge):
+        result_graph = transform_context_to_graph(results)
+
+    return {
+        "result": result_graph or results[0] if results and len(results) == 1 else results,
+        "graphs": graphs,
+        "context": context_texts,
+        "datasets": datasets,
+    }

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/modules/search/utils/transform_context_to_graph.py

@@ -14,7 +14,7 @@ def transform_context_to_graph(context: List[Edge]):
             if "name" in triplet.node1.attributes
             else triplet.node1.id,
             "type": triplet.node1.attributes["type"],
-            "attributes": triplet.node2.attributes,
+            "attributes": triplet.node1.attributes,
         }
         nodes[triplet.node2.id] = {
             "id": triplet.node2.id,

cognee-0.3.4.dev0/cognee/modules/search/utils/transform_insights_to_graph.py

@@ -0,0 +1,28 @@
+from typing import Dict, List, Tuple
+
+
+def transform_insights_to_graph(context: List[Tuple[Dict, Dict, Dict]]):
+    nodes = {}
+    edges = {}
+
+    for triplet in context:
+        nodes[triplet[0]["id"]] = {
+            "id": triplet[0]["id"],
+            "label": triplet[0]["name"] if "name" in triplet[0] else triplet[0]["id"],
+            "type": triplet[0]["type"],
+        }
+        nodes[triplet[2]["id"]] = {
+            "id": triplet[2]["id"],
+            "label": triplet[2]["name"] if "name" in triplet[2] else triplet[2]["id"],
+            "type": triplet[2]["type"],
+        }
+        edges[f"{triplet[0]['id']}_{triplet[1]['relationship_name']}_{triplet[2]['id']}"] = {
+            "source": triplet[0]["id"],
+            "target": triplet[2]["id"],
+            "label": triplet[1]["relationship_name"],
+        }
+
+    return {
+        "nodes": list(nodes.values()),
+        "edges": list(edges.values()),
+    }

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/tasks/temporal_graph/models.py

@@ -3,12 +3,17 @@ from pydantic import BaseModel, Field
 
 
 class Timestamp(BaseModel):
-    year: int = Field(..., ge=1, le=9999)
-    month: int = Field(..., ge=1, le=12)
-    day: int = Field(..., ge=1, le=31)
-    hour: int = Field(..., ge=0, le=23)
-    minute: int = Field(..., ge=0, le=59)
-    second: int = Field(..., ge=0, le=59)
+    year: int = Field(
+        ...,
+        ge=1,
+        le=9999,
+        description="Always required. If only a year is known, use it.",
+    )
+    month: int = Field(1, ge=1, le=12, description="If unknown, default to 1")
+    day: int = Field(1, ge=1, le=31, description="If unknown, default to 1")
+    hour: int = Field(0, ge=0, le=23, description="If unknown, default to 0")
+    minute: int = Field(0, ge=0, le=59, description="If unknown, default to 0")
+    second: int = Field(0, ge=0, le=59, description="If unknown, default to 0")
 
 
 class Interval(BaseModel):

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/tests/cli_tests/cli_unit_tests/test_cli_main.py

@@ -49,7 +49,7 @@ class TestCliMain:
     def test_main_no_command(self, mock_create_parser):
         """Test main function when no command is provided"""
         mock_parser = MagicMock()
-        mock_parser.parse_args.return_value = MagicMock(command=None)
+        mock_parser.parse_args.return_value = MagicMock(command=None, spec={})
         mock_create_parser.return_value = (mock_parser, {})
 
         result = main()
@@ -64,7 +64,7 @@ class TestCliMain:
         mock_command.execute.return_value = None
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -84,7 +84,7 @@ class TestCliMain:
         mock_command.execute.side_effect = CliCommandException("Test error", error_code=2)
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -103,7 +103,7 @@ class TestCliMain:
         mock_command.execute.side_effect = Exception("Generic error")
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -126,7 +126,7 @@ class TestCliMain:
         mock_command.execute.side_effect = test_exception
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})

{cognee-0.3.3 → cognee-0.3.4.dev0}/cognee/tests/test_temporal_graph.py

@@ -97,7 +97,7 @@ async def main():
         f"Expected exactly one DocumentChunk, but found {type_counts.get('DocumentChunk', 0)}"
     )
 
-    assert type_counts.get("Entity", 0) >= 20, (
+    assert type_counts.get("Entity", 0) >= 10, (
         f"Expected multiple entities (assert is set to 20), but found {type_counts.get('Entity', 0)}"
     )
 
@@ -105,52 +105,24 @@ async def main():
         f"Expected multiple entity types, but found {type_counts.get('EntityType', 0)}"
     )
 
-    assert type_counts.get("Event", 0) >= 20, (
+    assert type_counts.get("Event", 0) >= 10, (
         f"Expected multiple events (assert is set to 20), but found {type_counts.get('Event', 0)}"
     )
 
-    assert type_counts.get("Timestamp", 0) >= 20, (
-        f"Expected multiple timestamps (assert is set to 20), but found {type_counts.get('Timestamp', 0)}"
+    assert type_counts.get("Timestamp", 0) >= 10, (
+        f"Expected multiple timestamps (assert is set to 10), but found {type_counts.get('Timestamp', 0)}"
     )
 
-    assert type_counts.get("Interval", 0) >= 2, (
-        f"Expected multiple intervals, but found {type_counts.get('Interval', 0)}"
-    )
-
-    assert edge_type_counts.get("contains", 0) >= 20, (
+    assert edge_type_counts.get("contains", 0) >= 10, (
         f"Expected multiple 'contains' edge, but found {edge_type_counts.get('contains', 0)}"
     )
 
-    assert edge_type_counts.get("is_a", 0) >= 20, (
+    assert edge_type_counts.get("is_a", 0) >= 10, (
         f"Expected multiple 'is_a' edge, but found {edge_type_counts.get('is_a', 0)}"
     )
 
-    assert edge_type_counts.get("during", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of during and interval objects in the graph"
-    )
-
-    assert edge_type_counts.get("during", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of during and interval objects in the graph"
-    )
-
-    assert edge_type_counts.get("time_from", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of time_from and interval objects in the graph"
-    )
-
-    assert edge_type_counts.get("time_to", 0) == type_counts.get("Interval", 0), (
-        "Expected the same amount of time_to and interval objects in the graph"
-    )
-
     retriever = TemporalRetriever()
 
-    result_before = await retriever.extract_time_from_query("What happened before 1890?")
-
-    assert result_before[0] is None
-
-    result_after = await retriever.extract_time_from_query("What happened after 1891?")
-
-    assert result_after[1] is None
-
     result_between = await retriever.extract_time_from_query("What happened between 1890 and 1900?")
 
     assert result_between[1]