Merge pull request #176 from topoteretes/fix/integration-test-warnings

COG-485 - Fix/integration test warnings

cognee/api/client.py

@@ -146,7 +146,7 @@ class DatasetDTO(OutDTO):
     id: UUID
     name: str
     created_at: datetime
-    updated_at: Optional[datetime]
+    updated_at: Optional[datetime] = None
     owner_id: UUID
 
 @app.get("/api/v1/datasets", response_model = list[DatasetDTO])
@@ -200,7 +200,7 @@ class DataDTO(OutDTO):
     id: UUID
     name: str
     created_at: datetime
-    updated_at: Optional[datetime]
+    updated_at: Optional[datetime] = None
     extension: str
     mime_type: str
     raw_data_location: str

cognee/infrastructure/files/utils/extract_text_from_file.py

@@ -1,7 +1,8 @@
 from typing import BinaryIO
 from pypdf import PdfReader
+import filetype
 
-def extract_text_from_file(file: BinaryIO, file_type) -> str:
+def extract_text_from_file(file: BinaryIO, file_type: filetype.Type) -> str:
     """Extract text from a file"""
     if file_type.extension == "pdf":
         reader = PdfReader(stream = file)

cognee/modules/cognify/evaluate.py

@@ -1,4 +1,3 @@
-import dsp
 import dspy
 from dspy.evaluate.evaluate import Evaluate
 from dspy.primitives.example import Example

cognee/modules/cognify/train.py

@@ -1,4 +1,3 @@
-import dsp
 import dspy
 from dspy.teleprompt import BootstrapFewShot
 from dspy.primitives.example import Example

cognee/modules/pipelines/Pipeline.py

@@ -5,7 +5,7 @@
 
 class PipelineConfig(BaseModel):
     batch_count: int = 10
-    description: Optional[str]
+    description: Optional[str] = None
 
 class Pipeline():
     id: UUID = uuid4()

cognee/modules/pipelines/operations/run_parallel.py

@@ -1,8 +1,8 @@
-from typing import Any, Callable, Generator
+from typing import Any, Callable, Generator, List
 import asyncio
 from ..tasks.Task import Task
 
-def run_tasks_parallel(tasks: [Task]) -> Callable[[Any], Generator[Any, Any, Any]]:
+def run_tasks_parallel(tasks: List[Task]) -> Callable[[Any], Generator[Any, Any, Any]]:
     async def parallel_run(*args, **kwargs):
         parallel_tasks = [asyncio.create_task(task.run(*args, **kwargs)) for task in tasks]
 

cognee/shared/GithubTopology.py

@@ -18,7 +18,7 @@ class Directory(BaseModel):
     directories: List['Directory'] = []
 
 # Allows recursive Directory Model
-Directory.update_forward_refs()
+Directory.model_rebuild()
 
 class RepositoryProperties(BaseModel):
     custom_properties: Optional[Dict[str, Any]] = None

cognee/shared/SourceCodeGraph.py

@@ -6,15 +6,15 @@ class BaseClass(BaseModel):
     name: str
     type: Literal["Class"] = "Class"
     description: str
-    constructor_parameters: Optional[List[str]]
+    constructor_parameters: Optional[List[str]] = None
 
 class Class(BaseModel):
     id: str
     name: str
     type: Literal["Class"] = "Class"
     description: str
-    constructor_parameters: Optional[List[str]]
-    from_class: Optional[BaseClass]
+    constructor_parameters: Optional[List[str]] = None
+    from_class: Optional[BaseClass] = None
 
 class ClassInstance(BaseModel):
     id: str
@@ -28,7 +28,7 @@ class Function(BaseModel):
     name: str
     type: Literal["Function"] = "Function"
     description: str
-    parameters: Optional[List[str]]
+    parameters: Optional[List[str]] = None
     return_type: str
     is_static: Optional[bool] = False
 
@@ -38,7 +38,7 @@ class Variable(BaseModel):
     type: Literal["Variable"] = "Variable"
     description: str
     is_static: Optional[bool] = False
-    default_value: Optional[str]
+    default_value: Optional[str] = None
 
 class Operator(BaseModel):
     id: str

cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py

@@ -59,7 +59,7 @@ class Keyword(BaseModel):
             data_points.append(
                 DataPoint[Keyword](
                     id=str(classification_type_id),
-                    payload=Keyword.parse_obj({
+                    payload=Keyword.model_validate({
                         "uuid": str(classification_type_id),
                         "text": classification_type_label,
                         "chunk_id": str(data_chunk.chunk_id),
@@ -98,7 +98,7 @@ class Keyword(BaseModel):
                 data_points.append(
                     DataPoint[Keyword](
                         id=str(classification_subtype_id),
-                        payload=Keyword.parse_obj({
+                        payload=Keyword.model_validate({
                             "uuid": str(classification_subtype_id),
                             "text": classification_subtype_label,
                             "chunk_id": str(data_chunk.chunk_id),

cognee/tasks/infer_data_ontology/infer_data_ontology.py

@@ -56,7 +56,7 @@ async def recursive_flatten(self, items: Union[List[Dict[str, Any]], Dict[str, A
             for item in items:
                 flat_list.extend(await self.recursive_flatten(item, parent_id))
         elif isinstance(items, dict):
-            model = NodeModel.parse_obj(items)
+            model = NodeModel.model_validate(items)
             flat_list.append(await self.flatten_model(model, parent_id))
             for child in model.children:
                 flat_list.extend(await self.recursive_flatten(child, model.node_id))

cognee/tasks/infer_data_ontology/models/models.py

@@ -12,7 +12,7 @@ class NodeModel(BaseModel):
     default_relationship: Optional[RelationshipModel] = None
     children: List[Union[Dict[str, Any], "NodeModel"]] = Field(default_factory=list)
 
-NodeModel.update_forward_refs()
+NodeModel.model_rebuild()
 
 
 class OntologyNode(BaseModel):

cognee/tasks/save_chunks_to_store/save_chunks_to_store.py

@@ -11,7 +11,7 @@ async def save_chunks_to_store(data_chunks: list[DocumentChunk], collection_name
 
     # Remove and unlink existing chunks
     if await vector_engine.has_collection(collection_name):
-        existing_chunks = [DocumentChunk.parse_obj(chunk.payload) for chunk in (await vector_engine.retrieve(
+        existing_chunks = [DocumentChunk.model_validate(chunk.payload) for chunk in (await vector_engine.retrieve(
             collection_name,
             [str(chunk.chunk_id) for chunk in data_chunks],
         ))]

cognee/tests/test_library.py

@@ -49,7 +49,7 @@ async def  main():
 
     search_results = await cognee.search(SearchType.SUMMARIES, query = random_node_name)
     assert len(search_results) != 0, "Query related summaries don't exist."
-    print("\n\Extracted summaries are:\n")
+    print("\nExtracted summaries are:\n")
     for result in search_results:
         print(f"{result}\n")
 

cognee/tests/test_neo4j.py

@@ -53,7 +53,7 @@ async def main():
 
     search_results = await cognee.search(SearchType.SUMMARIES, query = random_node_name)
     assert len(search_results) != 0, "Query related summaries don't exist."
-    print("\n\Extracted summaries are:\n")
+    print("\nExtracted summaries are:\n")
     for result in search_results:
         print(f"{result}\n")
 

cognee/tests/test_qdrant.py

@@ -54,7 +54,7 @@ async def main():
 
     search_results = await cognee.search(SearchType.SUMMARIES, query = random_node_name)
     assert len(search_results) != 0, "Query related summaries don't exist."
-    print("\n\Extracted summaries are:\n")
+    print("\nExtracted summaries are:\n")
     for result in search_results:
         print(f"{result}\n")
 

cognee/tests/test_weaviate.py

@@ -52,7 +52,7 @@ async def main():
 
     search_results = await cognee.search(SearchType.SUMMARIES, query = random_node_name)
     assert len(search_results) != 0, "Query related summaries don't exist."
-    print("\n\Extracted summaries are:\n")
+    print("\nExtracted summaries are:\n")
     for result in search_results:
         print(f"{result}\n")