diff --git a/cognee/api/v1/permissions/routers/get_permissions_router.py b/cognee/api/v1/permissions/routers/get_permissions_router.py index d3217325b..78b822085 100644 --- a/cognee/api/v1/permissions/routers/get_permissions_router.py +++ b/cognee/api/v1/permissions/routers/get_permissions_router.py @@ -47,7 +47,7 @@ def get_permissions_router() -> APIRouter: group_id=group.id, permission_id=permission_entity.id ) ) - except IntegrityError as e: + except IntegrityError: raise EntityAlreadyExistsError(message="Group permission already exists.") await db.session.commit() @@ -70,7 +70,7 @@ def get_permissions_router() -> APIRouter: # Add association directly to the association table stmt = insert(UserGroup).values(user_id=user_id, group_id=group_id) await db.session.execute(stmt) - except IntegrityError as e: + except IntegrityError: raise EntityAlreadyExistsError(message="User is already part of group.") await db.session.commit() diff --git a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py index 9b2488b70..27061ce45 100644 --- a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py +++ b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py @@ -87,7 +87,7 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine): except litellm.exceptions.RateLimitError: if self.retry_count >= self.MAX_RETRIES: - raise Exception(f"Rate limit exceeded and no more retries left.") + raise Exception("Rate limit exceeded and no more retries left.") await exponential_backoff(self.retry_count) diff --git a/cognee/modules/users/get_user_db.py b/cognee/modules/users/get_user_db.py index c32aca79b..9eae0503e 100644 --- a/cognee/modules/users/get_user_db.py +++ b/cognee/modules/users/get_user_db.py @@ -1,5 +1,6 @@ from typing import AsyncGenerator from fastapi import Depends +from contextlib import asynccontextmanager from sqlalchemy.ext.asyncio import AsyncSession from fastapi_users.db import SQLAlchemyUserDatabase from cognee.infrastructure.databases.relational import get_relational_engine @@ -16,6 +17,6 @@ async def get_user_db(session: AsyncSession = Depends(get_async_session)): yield SQLAlchemyUserDatabase(session, User) -from contextlib import asynccontextmanager + get_user_db_context = asynccontextmanager(get_user_db) diff --git a/cognee/shared/utils.py b/cognee/shared/utils.py index 2d4271143..f257d450c 100644 --- a/cognee/shared/utils.py +++ b/cognee/shared/utils.py @@ -8,17 +8,19 @@ import hashlib from datetime import datetime, timezone import graphistry import networkx as nx -import numpy as np import pandas as pd import matplotlib.pyplot as plt import tiktoken import nltk import base64 -import networkx as nx -from bokeh.io import output_file, save from bokeh.plotting import figure, from_networkx from bokeh.models import Circle, MultiLine, HoverTool, ColumnDataSource, Range1d +from bokeh.plotting import output_file, show +from bokeh.embed import file_html +from bokeh.resources import CDN +import cairosvg +import logging from cognee.base_config import get_base_config from cognee.infrastructure.databases.graph import get_graph_engine @@ -272,16 +274,7 @@ def extract_pos_tags(sentence): return pos_tags -import networkx as nx -from bokeh.plotting import figure, output_file, show -from bokeh.models import Circle, MultiLine, HoverTool, Range1d -from bokeh.io import output_notebook -from bokeh.embed import file_html -from bokeh.resources import CDN -from bokeh.plotting import figure, from_networkx -import base64 -import cairosvg -import logging + logging.basicConfig(level=logging.INFO) diff --git a/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py b/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py index 7544bd398..36fa71c25 100644 --- a/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py +++ b/cognee/tests/unit/interfaces/graph/get_graph_from_model_test.py @@ -49,7 +49,7 @@ async def get_graph_from_model_test(): for document_chunk in document_chunks: document_chunk.contains.append( Entity( - name=f"Entity", + name="Entity", is_type=EntityType( name="Type 1", ), diff --git a/cognee/tests/unit/processing/utils/utils_test.py b/cognee/tests/unit/processing/utils/utils_test.py index e075db3b4..3abd46f38 100644 --- a/cognee/tests/unit/processing/utils/utils_test.py +++ b/cognee/tests/unit/processing/utils/utils_test.py @@ -102,8 +102,6 @@ def test_prepare_nodes(): assert len(nodes_df) == 1 -from unittest.mock import DEFAULT - def test_create_cognee_style_network_with_logo(): import networkx as nx diff --git a/evals/simple_rag_vs_cognee_eval.py b/evals/simple_rag_vs_cognee_eval.py index 9b7bd6a51..823099c26 100644 --- a/evals/simple_rag_vs_cognee_eval.py +++ b/evals/simple_rag_vs_cognee_eval.py @@ -1,14 +1,19 @@ from deepeval.dataset import EvaluationDataset from pydantic import BaseModel - +import os from typing import List, Type from deepeval.test_case import LLMTestCase import dotenv +from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.infrastructure.databases.vector import get_vector_engine +from cognee.base_config import get_base_config +import logging + +logger = logging.getLogger(__name__) dotenv.load_dotenv() -from cognee.infrastructure.llm.get_llm_client import get_llm_client dataset = EvaluationDataset() dataset.add_test_cases_from_json_file( @@ -39,9 +44,6 @@ print(dataset.goldens) print(dataset) -import logging - -logger = logging.getLogger(__name__) class AnswerModel(BaseModel): @@ -78,9 +80,6 @@ async def run_cognify_base_rag(): pass -import os -from cognee.base_config import get_base_config -from cognee.infrastructure.databases.vector import get_vector_engine async def cognify_search_base_rag(content: str, context: str):