Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

COG-485 - Fix/integration test warnings #176

Merged
merged 6 commits into from
Nov 5, 2024
Merged

COG-485 - Fix/integration test warnings #176

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
287427a
Update litellm and openai version, bump pydantic
Nov 4, 2024
79642cf
Replace update_forward_refs with model_rebuild calls
0xideas Nov 4, 2024
115e93b
Remove fastembed dependency
0xideas Nov 4, 2024
cb20506
Replace parse_obj with model_validate
0xideas Nov 4, 2024
5b07a0a
Add type hint in extract_text_from_file
0xideas Nov 4, 2024
67f5d07
Delete get_file_size
0xideas Nov 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions cognee/api/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down Expand Up @@ -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
Expand Down
25 changes: 0 additions & 25 deletions cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py
View file
Open in desktop

This file was deleted.

3 changes: 2 additions & 1 deletion cognee/infrastructure/files/utils/extract_text_from_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)
Expand Down
5 changes: 0 additions & 5 deletions cognee/infrastructure/files/utils/get_file_size.py
View file
Open in desktop

This file was deleted.

1 change: 0 additions & 1 deletion cognee/modules/cognify/evaluate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import dsp
import dspy
from dspy.evaluate.evaluate import Evaluate
from dspy.primitives.example import Example
Expand Down
1 change: 0 additions & 1 deletion cognee/modules/cognify/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import dsp
import dspy
from dspy.teleprompt import BootstrapFewShot
from dspy.primitives.example import Example
Expand Down
2 changes: 1 addition & 1 deletion cognee/modules/pipelines/Pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

class PipelineConfig(BaseModel):
batch_count: int = 10
description: Optional[str]
description: Optional[str] = None

class Pipeline():
id: UUID = uuid4()
Expand Down
4 changes: 2 additions & 2 deletions cognee/modules/pipelines/operations/run_parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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]

Expand Down
2 changes: 1 addition & 1 deletion cognee/shared/GithubTopology.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ class Directory(BaseModel):
directories: List['Directory'] = []

# Allows recursive Directory Model
Directory.update_forward_refs()
Directory.model_rebuild()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codebase verification

Several Pydantic v1 methods need to be updated to v2 equivalents

The codebase shows multiple instances of deprecated Pydantic v1 methods that should be updated to their v2 equivalents:

  • parse_obj() → model_validate() in:
    • cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py
    • cognee/tasks/infer_data_ontology/infer_data_ontology.py
    • cognee/tasks/save_chunks_to_store/save_chunks_to_store.py
  • schema related usages in database adapters should be reviewed and updated to use Pydantic v2's model configuration
  • The imports look good as they're using the new style with pydantic_settings for settings classes

While the change from update_forward_refs() to model_rebuild() is correct, there are other v1 methods that should also be updated to maintain consistency with Pydantic v2.

🔗 Analysis chain

LGTM! Correct update to use Pydantic v2's model_rebuild()

The change from update_forward_refs() to model_rebuild() aligns with Pydantic v2's best practices for handling forward references in recursive models.

Let's verify the Pydantic version and check for other potential v1 methods that might need updates:

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Description: Check Pydantic version and usage of deprecated v1 methods

# Test 1: Verify Pydantic version in dependencies
echo "Checking Pydantic version:"
grep -r "pydantic" pyproject.toml

# Test 2: Search for other potential v1 methods that might need updates
echo -e "\nChecking for other Pydantic v1 methods:"
rg -l "update_forward_refs|construct|validate|parse_obj|parse_raw|parse_file|schema|schema_json|validate_json" --type py

Length of output: 1287

Script:

#!/bin/bash
# Let's check the actual usage of these v1 methods in the files to determine if they need updates
echo "Checking specific v1 method usages:"
rg "update_forward_refs|construct|validate|parse_obj|parse_raw|parse_file|schema|schema_json|validate_json" --type py -C 2

# Also check for Pydantic imports to see if they're using v2 style imports
echo -e "\nChecking Pydantic import patterns:"
rg "from pydantic" --type py

Length of output: 27510


class RepositoryProperties(BaseModel):
custom_properties: Optional[Dict[str, Any]] = None
Expand Down
10 changes: 5 additions & 5 deletions cognee/shared/SourceCodeGraph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand Down
2 changes: 1 addition & 1 deletion cognee/tasks/infer_data_ontology/infer_data_ontology.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
2 changes: 1 addition & 1 deletion cognee/tasks/infer_data_ontology/models/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion cognee/tasks/save_chunks_to_store/save_chunks_to_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
))]
Expand Down
2 changes: 1 addition & 1 deletion cognee/tests/test_library.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
2 changes: 1 addition & 1 deletion cognee/tests/test_neo4j.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
2 changes: 1 addition & 1 deletion cognee/tests/test_qdrant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
2 changes: 1 addition & 1 deletion cognee/tests/test_weaviate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
Empty file added log.txt
View file
Open in desktop
Empty file.
Loading
Loading