From 860218632ffeee3b950f0c669f8c04285d33ac49 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Tue, 28 Jan 2025 17:15:25 +0100 Subject: [PATCH] refactor: add suggestions from PR Add suggestsions made by CodeRabbit on pull request --- cognee/infrastructure/llm/get_llm_client.py | 11 +++++------ .../data/processing/document_types/Document.py | 2 +- cognee/tasks/repo_processor/get_source_code_chunks.py | 2 -- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/cognee/infrastructure/llm/get_llm_client.py b/cognee/infrastructure/llm/get_llm_client.py index 383355fd2..ede8bd330 100644 --- a/cognee/infrastructure/llm/get_llm_client.py +++ b/cognee/infrastructure/llm/get_llm_client.py @@ -22,13 +22,12 @@ def get_llm_client(): # Check if max_token value is defined in liteLLM for given model # if not use value from cognee configuration - from cognee.infrastructure.llm.utils import get_model_max_tokens + from cognee.infrastructure.llm.utils import ( + get_model_max_tokens, + ) # imported here to avoid circular imports - max_tokens = ( - get_model_max_tokens(llm_config.llm_model) - if get_model_max_tokens(llm_config.llm_model) - else llm_config.llm_max_tokens - ) + model_max_tokens = get_model_max_tokens(llm_config.llm_model) + max_tokens = model_max_tokens if model_max_tokens else llm_config.llm_max_tokens if provider == LLMProvider.OPENAI: if llm_config.llm_api_key is None: diff --git a/cognee/modules/data/processing/document_types/Document.py b/cognee/modules/data/processing/document_types/Document.py index 80ba8e428..76ff1e045 100644 --- a/cognee/modules/data/processing/document_types/Document.py +++ b/cognee/modules/data/processing/document_types/Document.py @@ -11,5 +11,5 @@ class Document(DataPoint): mime_type: str _metadata: dict = {"index_fields": ["name"], "type": "Document"} - def read(self, chunk_size: int, chunker=str) -> str: + def read(self, chunk_size: int, chunker=str, max_chunk_tokens: Optional[int] = None) -> str: pass diff --git a/cognee/tasks/repo_processor/get_source_code_chunks.py b/cognee/tasks/repo_processor/get_source_code_chunks.py index 358d5218d..ca1c76e46 100644 --- a/cognee/tasks/repo_processor/get_source_code_chunks.py +++ b/cognee/tasks/repo_processor/get_source_code_chunks.py @@ -97,8 +97,6 @@ def _get_chunk_source_code( current_source_code = "" # Get embedding engine used in vector database - from cognee.infrastructure.databases.vector.get_vector_engine import get_vector_engine - embedding_engine = get_vector_engine().embedding_engine for i, (child_code, token_count) in enumerate(code_token_counts):