Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
170 commits
Select commit Hold shift + click to select a range
af94012
added files
Vasilije1990 Jul 26, 2025
84ba0b2
fix cognify
Vasilije1990 Aug 4, 2025
1a15669
Merge branch 'dev' into add_cli
Vasilije1990 Aug 4, 2025
8590194
feat: add top_k input to SearchView and send to backend
Raj2604 Aug 5, 2025
e26c216
Merge branch 'dev' into feature/top_k-input
Raj2604 Aug 5, 2025
ff1614d
format
Vasilije1990 Aug 12, 2025
1de81c8
added fixes
Vasilije1990 Aug 12, 2025
b226595
added fix
Vasilije1990 Aug 12, 2025
e3bd474
git merge
Vasilije1990 Aug 12, 2025
a15cad0
Merge branch 'dev' into add_cli
Vasilije1990 Aug 12, 2025
ed555a7
feat: adds new Exception classes
hajdul88 Aug 13, 2025
0555dd9
Update __init__.py
hajdul88 Aug 13, 2025
1b3898d
feat: updates rel_db exceptions with the new error classes
hajdul88 Aug 13, 2025
b40dbf3
feat: adds new error classes to litellm instructor
hajdul88 Aug 13, 2025
623148e
feat: adds new errors to litellm client
hajdul88 Aug 13, 2025
7bd2660
feat: setting base classes of data exceptions to the new ones
hajdul88 Aug 13, 2025
6870bba
feat: adds new error to delete
hajdul88 Aug 13, 2025
ebd4403
adds new error classes to keywordextration error and defines new error
hajdul88 Aug 13, 2025
657c775
feat: adds configattributeerror
hajdul88 Aug 13, 2025
5bc00f1
feat: adds new search classes to search.py
hajdul88 Aug 13, 2025
9fb9f68
adds new base errors to retrieval exceptions
hajdul88 Aug 13, 2025
91b9c11
changes base class for vectordb exceptions
hajdul88 Aug 13, 2025
6dcd59c
feat: Adds changes to cognee graph part
hajdul88 Aug 13, 2025
32996aa
feat: adds new error classes to llm and databases + introduces loglev…
hajdul88 Aug 13, 2025
bed523a
feat: api error handling restruct
hajdul88 Aug 13, 2025
38329da
Merge branch 'dev' into feature/cog-2717-add-better-error-management-…
hajdul88 Aug 13, 2025
544e089
feat: removing invalidValueErrors
hajdul88 Aug 13, 2025
885f7c3
chore: fixing graph elements tests
hajdul88 Aug 13, 2025
d1bfeaa
fix: fixes search unit test error expectation
hajdul88 Aug 13, 2025
fabbd63
chore: renaming error
hajdul88 Aug 13, 2025
59c9204
Update exceptions.py
hajdul88 Aug 13, 2025
68327d3
chore: Changing base classes for exceptions that were already defined
hajdul88 Aug 13, 2025
da40365
ruff formatting
hajdul88 Aug 13, 2025
7af4775
chore: changed pypdf error base class
hajdul88 Aug 13, 2025
66d2c75
chore: changing docstring
hajdul88 Aug 13, 2025
8bd2416
chore: fixing delete exception import in dataset router
hajdul88 Aug 13, 2025
d14d31a
chore: updating delete_by_id test
hajdul88 Aug 13, 2025
748e9fa
feat: adds s3 file system not found error to ingestion
hajdul88 Aug 13, 2025
dbb967f
chore: updates neptune exception base + ruff
hajdul88 Aug 13, 2025
ca21122
Merge branch 'dev' into feature/cog-2717-add-better-error-management-…
hajdul88 Aug 14, 2025
c75f017
Merge branch 'dev' into feature/cog-2717-add-better-error-management-…
hajdul88 Aug 14, 2025
c99b453
feat: adds WrongDataDocumentError to classify documents
hajdul88 Aug 14, 2025
df3a3df
feat: adds errors to classify, and chunking top level
hajdul88 Aug 14, 2025
9f965c4
feat: adds input checks for extract graph from data
hajdul88 Aug 14, 2025
63d071f
feat: adds input checks for add datapoints and summarization tasks
hajdul88 Aug 14, 2025
affbc55
chore: ruff formatting
hajdul88 Aug 14, 2025
a7d7e12
ruff fix
hajdul88 Aug 14, 2025
dd6e26f
Merge remote-tracking branch 'origin/dev' into feature/cog-2717-add-b…
hajdul88 Aug 14, 2025
d884cc4
Merge branch 'dev' into feature/cog-2717-add-better-error-management-…
hajdul88 Aug 15, 2025
f5eecba
feat: add top_k-input and send it to backend
Raj2604 Aug 10, 2025
c606273
Refactor CI workflows to replace Poetry with uv for dependency manage…
daukadolt Aug 14, 2025
cb5590b
Update installation instructions and replace Poetry with uv across do…
daukadolt Aug 14, 2025
a1edfc9
docs: update README.md with new supported python versions
daukadolt Aug 14, 2025
1ab3328
fix: uv uninstalls rest of packages in some workflows
daukadolt Aug 14, 2025
bcdbadc
fix: unintentionally uninstall required deps when "uv sync"
daukadolt Aug 14, 2025
3941e46
chore: add step to rebuild uv lockfile in CI workflow
daukadolt Aug 15, 2025
4312508
undo: keep poetry-uv change limited to CI/CD fo now
daukadolt Aug 15, 2025
216ffd2
Refactor CI workflows to replace Poetry with uv (#1250)
Vasilije1990 Aug 15, 2025
81a1758
Remove NetworkX from cognee core.
misselvexu Aug 13, 2025
85c7f11
Keep networkx as core dependency
daukadolt Aug 15, 2025
50ebced
fix: natural_language_retriever.py
daukadolt Aug 15, 2025
8d045b9
chore: forgot to remove NetworkX check in NLPRetriever
daukadolt Aug 15, 2025
8825aaa
format: ruff format
daukadolt Aug 15, 2025
1ee697b
fix: Return distributed as part of Cognee build
dexters1 Aug 15, 2025
1e57ac9
fix: Return distributed as part of Cognee build (#1257)
Vasilije1990 Aug 15, 2025
271e7e9
feature: Introduces new error handling (4 base errors + specified hie…
Vasilije1990 Aug 15, 2025
86e5135
chore: Update Cognee version
dexters1 Aug 15, 2025
99a517d
chore: Update Cognee version (#1258)
dexters1 Aug 15, 2025
761bb1b
test: deprecate networkx_metrics_test.py
daukadolt Aug 15, 2025
5f7598d
test: use neo4j_metrics_test in descriptive tests instead of networkx
daukadolt Aug 15, 2025
e4e0512
feat: add reusable GitHub Action to set up Neo4j with Graph Data Scie…
daukadolt Aug 15, 2025
73e1aa7
Merge branch 'dev' into issues/1183
Vasilije1990 Aug 15, 2025
9987e05
Remove NetworkX from cognee core. (#1241)
Vasilije1990 Aug 15, 2025
b0e3f89
move to gpt5
Vasilije1990 Aug 17, 2025
c4ec679
Merge branch 'dev' into move_to_gpt5
Vasilije1990 Aug 17, 2025
1bd40f1
renamed max tokens
Vasilije1990 Aug 17, 2025
693c883
updated deepeval
Vasilije1990 Aug 17, 2025
abbce6e
fix: Move to gpt5 (#1262)
Vasilije1990 Aug 17, 2025
711c805
feat: adds cognee-user interactions to search
hajdul88 Aug 18, 2025
d53ebb2
Merge branch 'dev' into feature/cog-2734-cognee-feedbacks-interaction…
hajdul88 Aug 18, 2025
dc637f7
fix: fixes add datapoints params
hajdul88 Aug 18, 2025
fbb7d72
fix: ruff formatting
hajdul88 Aug 18, 2025
1d63da7
chore: removes duplicated func def
hajdul88 Aug 18, 2025
78fb415
chore: changes context return value in tests
hajdul88 Aug 18, 2025
b6be617
fix: fixes tests
hajdul88 Aug 18, 2025
0529d4b
fix: fixes kuzu and neo4j tests
hajdul88 Aug 18, 2025
9a46d14
chore: fix search db tests
hajdul88 Aug 18, 2025
fc43ac7
feat: adds user feedback search type
hajdul88 Aug 18, 2025
4d021ec
chore: fix openai version (#1266)
hajdul88 Aug 18, 2025
d4ace62
Merge branch 'dev' into feature/cog-2734-cognee-feedbacks-interaction…
hajdul88 Aug 18, 2025
0fbe218
chore: fixes ruff
hajdul88 Aug 18, 2025
fb51131
Merge branch 'dev' into add_cli
Vasilije1990 Aug 18, 2025
9017b63
simplify logging utils and refactor
Vasilije1990 Aug 18, 2025
9a17aa7
added fix
Vasilije1990 Aug 18, 2025
13e3e3b
added fix
Vasilije1990 Aug 18, 2025
cae173b
fix bug
Vasilije1990 Aug 18, 2025
d084d00
added tests
Vasilije1990 Aug 18, 2025
66d7fa9
fix tests
Vasilije1990 Aug 18, 2025
3152368
added fix
Vasilije1990 Aug 18, 2025
f5d702f
added fixes to integraton tests
Vasilije1990 Aug 19, 2025
fe6c900
added fix
Vasilije1990 Aug 19, 2025
372181d
fix: fixes unit test
hajdul88 Aug 19, 2025
fcdee16
feat: adds kuzu and neo4j tests for feedback and interaction features
hajdul88 Aug 19, 2025
4e31ae7
chore: deletes unused var from search test
hajdul88 Aug 19, 2025
e912b6f
fix: ensure connection check before executing checkpoint in KuzuAdapter
daukadolt Aug 19, 2025
c6ec22a
feat: adds scores to Feedback node
hajdul88 Aug 19, 2025
67c7919
fix: ensure connection check before executing checkpoint in KuzuAdapt…
Vasilije1990 Aug 19, 2025
bf1970b
feat: Add top_k Input Control to Search UI for Adjustable Graph Explo…
Vasilije1990 Aug 19, 2025
4a5d5f7
feat: adds feedback weights to edges
hajdul88 Aug 19, 2025
f5d8fc6
chore: ruff ruff
hajdul88 Aug 19, 2025
67d88c4
Merge branch 'dev' into feature/cog-2734-cognee-feedbacks-interaction…
hajdul88 Aug 19, 2025
b8cac4c
feat: adds weight test at the end of test_search_db
hajdul88 Aug 19, 2025
6d9a100
feature: Introduces Cognee-user interactions feature and feedback sea…
Vasilije1990 Aug 19, 2025
d69669b
added ability to send custom prompts to cognify
Vasilije1990 Aug 22, 2025
6d438c8
feat: Add dataset and pipeline status layer
dexters1 Aug 25, 2025
6f230c5
feature: adds environment_setup_and_checks general purpose layer to c…
hajdul88 Aug 25, 2025
a805c64
Fix, remove, improve cli tests.
siillee Aug 25, 2025
573e24f
Remove duplicate test files.
siillee Aug 25, 2025
d4b23aa
feat: Add process_pipeline_check
dexters1 Aug 25, 2025
9a4e8dd
refactor: Make the authorized_user_datasets function more understandable
dexters1 Aug 25, 2025
e85fd35
Fix linter error.
siillee Aug 25, 2025
a3d5164
Fix formatting. Potential fix to invalid api key error.
siillee Aug 25, 2025
66b8bb3
Another potential fix for the invalid API key error.
siillee Aug 25, 2025
31b33a9
Yet another potential fix for the invalid API key.
siillee Aug 25, 2025
4b593aa
And another potential fix for the API error.
siillee Aug 25, 2025
454358e
Solution for the API key error.
siillee Aug 25, 2025
de7a5a1
Fixed yml files
siillee Aug 25, 2025
dce525e
Fix again.
siillee Aug 25, 2025
cf7d41c
Fix
siillee Aug 26, 2025
a6a33e9
Feature/cog 2768 inspect cli tests fix failing ones (#1286)
Vasilije1990 Aug 26, 2025
950d29a
refactor: Update typing
dexters1 Aug 26, 2025
03fc3c8
Merge branch 'dev' into cognee-pipeline-layers
dexters1 Aug 26, 2025
d91b0f6
feature: adds pipeline execution layer to cognify (#1291)
hajdul88 Aug 26, 2025
42d33fc
fix: fixes search test behaviour and adds comments to new pipeline ex…
hajdul88 Aug 26, 2025
5ea3056
refactor: Rename auth layer
dexters1 Aug 26, 2025
0d77649
Merge branch 'cognee-pipeline-layers' of github.com:topoteretes/cogne…
dexters1 Aug 26, 2025
5771b36
feat: cognee pipeline layers (#1287)
dexters1 Aug 26, 2025
ea4f58e
feat: migrate pipeline input validation to a layer (#1284)
borisarzentar Aug 26, 2025
c013499
Fix tests failing for python 3.10
siillee Aug 26, 2025
007399b
feat: migrate pipeline status reset to add pipeline (#1289)
borisarzentar Aug 26, 2025
e5f0c22
Merge branch 'dev' into add_cli
Vasilije1990 Aug 26, 2025
d6e6e87
feat: Add cli (#1197)
Vasilije1990 Aug 26, 2025
8c69653
fix: Resolve issue with Windows path
dexters1 Aug 26, 2025
65542ec
refactor: Make CI/CD faster add more OS tests
dexters1 Aug 26, 2025
229a7a1
refactor: Speed up CI/CD execution time
dexters1 Aug 26, 2025
9c31617
fix: Resovle windows file path issue
dexters1 Aug 26, 2025
58655ca
refactor: Add proper path to test file
dexters1 Aug 26, 2025
6e5acec
refactor: make run_pipeline a high-level api for running pipelines (#…
borisarzentar Aug 27, 2025
23a46e5
Merge branch 'dev' into fix-windows-path
dexters1 Aug 27, 2025
ecc403e
fix: Add binary for python magic for windows os
dexters1 Aug 27, 2025
644116a
Merge branch 'fix-windows-path' of github.com:topoteretes/cognee into…
dexters1 Aug 27, 2025
eb65a89
fix: Resolve integration tests path issue
dexters1 Aug 27, 2025
23ea1c1
Potential fix for code scanning alert no. 187: Workflow does not cont…
dexters1 Aug 27, 2025
e4e1a54
refactor: Add read permissions only for gh token
dexters1 Aug 27, 2025
11934cd
fix: Resolve integration test issue
dexters1 Aug 27, 2025
3a3274b
Update .github/workflows/test_different_operating_systems.yml
dexters1 Aug 27, 2025
24c155b
fix: Resolve issue with Windows path (#1295)
Vasilije1990 Aug 27, 2025
62afced
feat: Added custom prompt to cognify (#1278)
Vasilije1990 Aug 27, 2025
d4ce0e1
Merge branch 'main' into merge-main-vol-6
dexters1 Aug 27, 2025
6ca46f1
refactor: ruff format
dexters1 Aug 27, 2025
328b75a
Merge main vol 6 (#1298)
Vasilije1990 Aug 27, 2025
5e6ada2
fix issue
Vasilije1990 Aug 27, 2025
792452b
added fix
Vasilije1990 Aug 27, 2025
75ea520
fix to cli
Vasilije1990 Aug 27, 2025
c933290
fixes to CLI
Vasilije1990 Aug 27, 2025
9df3466
format
Vasilije1990 Aug 27, 2025
19b59fc
Update Jupyter notebooks: added execution counts, improved logging ou…
daukadolt Aug 27, 2025
9f27402
Update Jupyter notebooks (#1299)
Vasilije1990 Aug 27, 2025
235015b
linter fix
Vasilije1990 Aug 27, 2025
cc5e684
added fix to tests
Vasilije1990 Aug 27, 2025
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
Prev Previous commit
Next Next commit
renamed max tokens
  • Loading branch information
Vasilije1990 committed Aug 17, 2025
commit 1bd40f1401ffcfb4c1ca4256ecaef45e76ad1383
2 changes: 1 addition & 1 deletion cognee/api/v1/cognify/cognify.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ async def cognify(
- LangchainChunker: Recursive character splitting with overlap
Determines how documents are segmented for processing.
chunk_size: Maximum tokens per chunk. Auto-calculated based on LLM if None.
Formula: min(embedding_max_tokens, llm_max_tokens // 2)
Formula: min(embedding_max_completion_tokens, llm_max_completion_tokens // 2)
Default limits: ~512-8192 tokens depending on models.
Smaller chunks = more granular but potentially fragmented knowledge.
ontology_file_path: Path to RDF/OWL ontology file for domain-specific entity types.
Expand Down
2 changes: 1 addition & 1 deletion cognee/api/v1/responses/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ class ResponseRequest(InDTO):
tool_choice: Optional[Union[str, Dict[str, Any]]] = "auto"
user: Optional[str] = None
temperature: Optional[float] = 1.0
max_tokens: Optional[int] = None
max_completion_tokens: Optional[int] = None


class ToolCallOutput(BaseModel):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,11 +41,11 @@ def __init__(
self,
model: Optional[str] = "openai/text-embedding-3-large",
dimensions: Optional[int] = 3072,
max_tokens: int = 512,
max_completion_tokens: int = 512,
):
self.model = model
self.dimensions = dimensions
self.max_tokens = max_tokens
self.max_completion_tokens = max_completion_tokens
self.tokenizer = self.get_tokenizer()
# self.retry_count = 0
self.embedding_model = TextEmbedding(model_name=model)
Expand Down Expand Up @@ -112,7 +112,9 @@ def get_tokenizer(self):
"""
logger.debug("Loading tokenizer for FastembedEmbeddingEngine...")

tokenizer = TikTokenTokenizer(model="gpt-4o", max_tokens=self.max_tokens)
tokenizer = TikTokenTokenizer(
model="gpt-4o", max_completion_tokens=self.max_completion_tokens
)

logger.debug("Tokenizer loaded for for FastembedEmbeddingEngine")
return tokenizer
Original file line number Diff line number Diff line change
Expand Up @@ -57,15 +57,15 @@ def __init__(
api_key: str = None,
endpoint: str = None,
api_version: str = None,
max_tokens: int = 512,
max_completion_tokens: int = 512,
):
self.api_key = api_key
self.endpoint = endpoint
self.api_version = api_version
self.provider = provider
self.model = model
self.dimensions = dimensions
self.max_tokens = max_tokens
self.max_completion_tokens = max_completion_tokens
self.tokenizer = self.get_tokenizer()
self.retry_count = 0

Expand Down Expand Up @@ -179,20 +179,29 @@ def get_tokenizer(self):
model = self.model.split("/")[-1]

if "openai" in self.provider.lower():
tokenizer = TikTokenTokenizer(model=model, max_tokens=self.max_tokens)
tokenizer = TikTokenTokenizer(
model=model, max_completion_tokens=self.max_completion_tokens
)
elif "gemini" in self.provider.lower():
tokenizer = GeminiTokenizer(model=model, max_tokens=self.max_tokens)
tokenizer = GeminiTokenizer(
model=model, max_completion_tokens=self.max_completion_tokens
)
elif "mistral" in self.provider.lower():
tokenizer = MistralTokenizer(model=model, max_tokens=self.max_tokens)
tokenizer = MistralTokenizer(
model=model, max_completion_tokens=self.max_completion_tokens
)
else:
try:
tokenizer = HuggingFaceTokenizer(
model=self.model.replace("hosted_vllm/", ""), max_tokens=self.max_tokens
model=self.model.replace("hosted_vllm/", ""),
max_completion_tokens=self.max_completion_tokens,
)
except Exception as e:
logger.warning(f"Could not get tokenizer from HuggingFace due to: {e}")
logger.info("Switching to TikToken default tokenizer.")
tokenizer = TikTokenTokenizer(model=None, max_tokens=self.max_tokens)
tokenizer = TikTokenTokenizer(
model=None, max_completion_tokens=self.max_completion_tokens
)

logger.debug(f"Tokenizer loaded for model: {self.model}")
return tokenizer
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
Instance variables:
- model
- dimensions
- max_tokens
- max_completion_tokens
- endpoint
- mock
- huggingface_tokenizer_name
Expand All @@ -39,7 +39,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):

model: str
dimensions: int
max_tokens: int
max_completion_tokens: int
endpoint: str
mock: bool
huggingface_tokenizer_name: str
Expand All @@ -50,13 +50,13 @@ def __init__(
self,
model: Optional[str] = "avr/sfr-embedding-mistral:latest",
dimensions: Optional[int] = 1024,
max_tokens: int = 512,
max_completion_tokens: int = 512,
endpoint: Optional[str] = "http://localhost:11434/api/embeddings",
huggingface_tokenizer: str = "Salesforce/SFR-Embedding-Mistral",
):
self.model = model
self.dimensions = dimensions
self.max_tokens = max_tokens
self.max_completion_tokens = max_completion_tokens
self.endpoint = endpoint
self.huggingface_tokenizer_name = huggingface_tokenizer
self.tokenizer = self.get_tokenizer()
Expand Down Expand Up @@ -132,7 +132,7 @@ def get_tokenizer(self):
"""
logger.debug("Loading HuggingfaceTokenizer for OllamaEmbeddingEngine...")
tokenizer = HuggingFaceTokenizer(
model=self.huggingface_tokenizer_name, max_tokens=self.max_tokens
model=self.huggingface_tokenizer_name, max_completion_tokens=self.max_completion_tokens
)
logger.debug("Tokenizer loaded for OllamaEmbeddingEngine")
return tokenizer
4 changes: 2 additions & 2 deletions cognee/infrastructure/databases/vector/embeddings/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ class EmbeddingConfig(BaseSettings):
embedding_endpoint: Optional[str] = None
embedding_api_key: Optional[str] = None
embedding_api_version: Optional[str] = None
embedding_max_tokens: Optional[int] = 8191
embedding_max_completion_tokens: Optional[int] = 8191
Copy link
Contributor

Choose a reason for hiding this comment

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

🛠️ Refactor suggestion

Backwards compatibility for env var rename.

Renaming to embedding_max_completion_tokens will break existing .env using EMBEDDING_MAX_TOKENS. Accept both via validation alias.

-from typing import Optional
+from typing import Optional
+from pydantic import Field, AliasChoices
@@
-    embedding_max_completion_tokens: Optional[int] = 8191
+    embedding_max_completion_tokens: Optional[int] = Field(
+        8191,
+        validation_alias=AliasChoices(
+            "embedding_max_completion_tokens",
+            "embedding_max_tokens",
+            "EMBEDDING_MAX_COMPLETION_TOKENS",
+            "EMBEDDING_MAX_TOKENS",
+        ),
+    )
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
embedding_max_completion_tokens: Optional[int] = 8191
from typing import Optional
from pydantic import Field, AliasChoices
embedding_max_completion_tokens: Optional[int] = Field(
8191,
validation_alias=AliasChoices(
"embedding_max_completion_tokens",
"embedding_max_tokens",
"EMBEDDING_MAX_COMPLETION_TOKENS",
"EMBEDDING_MAX_TOKENS",
),
)

huggingface_tokenizer: Optional[str] = None
model_config = SettingsConfigDict(env_file=".env", extra="allow")

Expand All @@ -38,7 +38,7 @@ def to_dict(self) -> dict:
"embedding_endpoint": self.embedding_endpoint,
"embedding_api_key": self.embedding_api_key,
"embedding_api_version": self.embedding_api_version,
"embedding_max_tokens": self.embedding_max_tokens,
"embedding_max_completion_tokens": self.embedding_max_completion_tokens,
"huggingface_tokenizer": self.huggingface_tokenizer,
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ def get_embedding_engine() -> EmbeddingEngine:
config.embedding_provider,
config.embedding_model,
config.embedding_dimensions,
config.embedding_max_tokens,
config.embedding_max_completion_tokens,
config.embedding_endpoint,
config.embedding_api_key,
config.embedding_api_version,
Expand All @@ -41,7 +41,7 @@ def create_embedding_engine(
embedding_provider,
embedding_model,
embedding_dimensions,
embedding_max_tokens,
embedding_max_completion_tokens,
embedding_endpoint,
embedding_api_key,
embedding_api_version,
Expand All @@ -58,7 +58,7 @@ def create_embedding_engine(
'ollama', or another supported provider.
- embedding_model: The model to be used for the embedding engine.
- embedding_dimensions: The number of dimensions for the embeddings.
- embedding_max_tokens: The maximum number of tokens for the embeddings.
- embedding_max_completion_tokens: The maximum number of tokens for the embeddings.
- embedding_endpoint: The endpoint for the embedding service, relevant for certain
providers.
- embedding_api_key: API key to authenticate with the embedding service, if
Expand All @@ -81,7 +81,7 @@ def create_embedding_engine(
return FastembedEmbeddingEngine(
model=embedding_model,
dimensions=embedding_dimensions,
max_tokens=embedding_max_tokens,
max_completion_tokens=embedding_max_completion_tokens,
)

if embedding_provider == "ollama":
Expand All @@ -90,7 +90,7 @@ def create_embedding_engine(
return OllamaEmbeddingEngine(
model=embedding_model,
dimensions=embedding_dimensions,
max_tokens=embedding_max_tokens,
max_completion_tokens=embedding_max_completion_tokens,
endpoint=embedding_endpoint,
huggingface_tokenizer=huggingface_tokenizer,
)
Expand All @@ -104,5 +104,5 @@ def create_embedding_engine(
api_version=embedding_api_version,
model=embedding_model,
dimensions=embedding_dimensions,
max_tokens=embedding_max_tokens,
max_completion_tokens=embedding_max_completion_tokens,
)
6 changes: 3 additions & 3 deletions cognee/infrastructure/llm/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ class LLMConfig(BaseSettings):
- llm_api_version
- llm_temperature
- llm_streaming
- llm_max_tokens
- llm_max_completion_tokens
- transcription_model
- graph_prompt_path
- llm_rate_limit_enabled
Expand All @@ -41,7 +41,7 @@ class LLMConfig(BaseSettings):
llm_api_version: Optional[str] = None
llm_temperature: float = 0.0
llm_streaming: bool = False
llm_max_tokens: int = 16384
llm_max_completion_tokens: int = 16384

baml_llm_provider: str = "openai"
baml_llm_model: str = "gpt-5-mini"
Expand Down Expand Up @@ -171,7 +171,7 @@ def to_dict(self) -> dict:
"api_version": self.llm_api_version,
"temperature": self.llm_temperature,
"streaming": self.llm_streaming,
"max_tokens": self.llm_max_tokens,
"max_completion_tokens": self.llm_max_completion_tokens,
"transcription_model": self.transcription_model,
"graph_prompt_path": self.graph_prompt_path,
"rate_limit_enabled": self.llm_rate_limit_enabled,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,15 @@ class AnthropicAdapter(LLMInterface):
name = "Anthropic"
model: str

def __init__(self, max_tokens: int, model: str = None):
def __init__(self, max_completion_tokens: int, model: str = None):
import anthropic

self.aclient = instructor.patch(
create=anthropic.AsyncAnthropic().messages.create, mode=instructor.Mode.ANTHROPIC_TOOLS
)

self.model = model
self.max_tokens = max_tokens
self.max_completion_tokens = max_completion_tokens

Comment on lines +26 to 35
Copy link
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue

Fix: hard-coded max_completion_tokens ignores constructor value

acreate_structured_output hard-codes 4096 and bypasses the configured max_completion_tokens. Use the instance attribute to honor caller configuration.

Apply:

@@
         return await self.aclient(
             model=self.model,
-            max_completion_tokens=4096,
+            max_completion_tokens=self.max_completion_tokens,
             max_retries=5,

Also applies to: 58-61

🤖 Prompt for AI Agents
In
cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py
around lines 26-35 (and similarly at lines 58-61), the acreate_structured_output
call is hard-coded to 4096 for max_completion_tokens which ignores the
constructor parameter; replace the hard-coded literal with the instance
attribute self.max_completion_tokens so the configured value passed to __init__
is used, ensuring both occurrences use self.max_completion_tokens instead of
4096 and keeping the rest of the call intact.

@sleep_and_retry_async()
@rate_limit_async
Expand All @@ -57,7 +57,7 @@ async def acreate_structured_output(

return await self.aclient(
model=self.model,
max_tokens=4096,
max_completion_tokens=4096,
max_retries=5,
messages=[
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ def __init__(
self,
api_key: str,
model: str,
max_tokens: int,
max_completion_tokens: int,
endpoint: Optional[str] = None,
api_version: Optional[str] = None,
streaming: bool = False,
Expand All @@ -44,7 +44,7 @@ def __init__(
self.endpoint = endpoint
self.api_version = api_version
self.streaming = streaming
self.max_tokens = max_tokens
self.max_completion_tokens = max_completion_tokens

@observe(as_type="generation")
@sleep_and_retry_async()
Expand Down Expand Up @@ -90,7 +90,7 @@ async def acreate_structured_output(
model=f"{self.model}",
messages=messages,
api_key=self.api_key,
max_tokens=self.max_tokens,
max_completion_tokens=self.max_completion_tokens,
temperature=0.1,
response_format=response_schema,
timeout=100,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ def __init__(
api_key: str,
model: str,
name: str,
max_tokens: int,
max_completion_tokens: int,
fallback_model: str = None,
fallback_api_key: str = None,
fallback_endpoint: str = None,
Expand All @@ -50,7 +50,7 @@ def __init__(
self.model = model
self.api_key = api_key
self.endpoint = endpoint
self.max_tokens = max_tokens
self.max_completion_tokens = max_completion_tokens

self.fallback_model = fallback_model
self.fallback_api_key = fallback_api_key
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,11 +54,15 @@ 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,
get_model_max_completion_tokens,
) # imported here to avoid circular imports

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
model_max_completion_tokens = get_model_max_completion_tokens(llm_config.llm_model)
max_completion_tokens = (
model_max_completion_tokens
if model_max_completion_tokens
else llm_config.llm_max_completion_tokens
)

if provider == LLMProvider.OPENAI:
if llm_config.llm_api_key is None:
Expand All @@ -74,7 +78,7 @@ def get_llm_client():
api_version=llm_config.llm_api_version,
model=llm_config.llm_model,
transcription_model=llm_config.transcription_model,
max_tokens=max_tokens,
max_completion_tokens=max_completion_tokens,
streaming=llm_config.llm_streaming,
fallback_api_key=llm_config.fallback_api_key,
fallback_endpoint=llm_config.fallback_endpoint,
Expand All @@ -94,15 +98,17 @@ def get_llm_client():
llm_config.llm_api_key,
llm_config.llm_model,
"Ollama",
max_tokens=max_tokens,
max_completion_tokens=max_completion_tokens,
)

elif provider == LLMProvider.ANTHROPIC:
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.anthropic.adapter import (
AnthropicAdapter,
)

return AnthropicAdapter(max_tokens=max_tokens, model=llm_config.llm_model)
return AnthropicAdapter(
max_completion_tokens=max_completion_tokens, model=llm_config.llm_model
)

elif provider == LLMProvider.CUSTOM:
if llm_config.llm_api_key is None:
Expand All @@ -117,7 +123,7 @@ def get_llm_client():
llm_config.llm_api_key,
llm_config.llm_model,
"Custom",
max_tokens=max_tokens,
max_completion_tokens=max_completion_tokens,
fallback_api_key=llm_config.fallback_api_key,
fallback_endpoint=llm_config.fallback_endpoint,
fallback_model=llm_config.fallback_model,
Expand All @@ -134,7 +140,7 @@ def get_llm_client():
return GeminiAdapter(
api_key=llm_config.llm_api_key,
model=llm_config.llm_model,
max_tokens=max_tokens,
max_completion_tokens=max_completion_tokens,
endpoint=llm_config.llm_endpoint,
api_version=llm_config.llm_api_version,
streaming=llm_config.llm_streaming,
Expand Down
Loading
Loading