gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: add variable to control instructor mode

Browse source
This commit is contained in:
Andrej Milicevic 2025-11-13 18:25:07 +01:00
parent a5bd504daa
commit 2337d36f7b
7 changed files with 58 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/infrastructure/llm/config.py
View file

@ -38,6 +38,7 @@ class LLMConfig(BaseSettings):
"""
structured_output_framework: str = "instructor"
llm_instructor_mode: Optional[str] = None
llm_provider: str = "openai"
llm_model: str = "openai/gpt-5-mini"
llm_endpoint: str = ""
@ -181,6 +182,7 @@ class LLMConfig(BaseSettings):
instance.
"""
return {
"llm_instructor_mode": self.llm_instructor_mode,
"provider": self.llm_provider,
"model": self.llm_model,
"endpoint": self.llm_endpoint,

8
cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py
View file

@ -28,13 +28,19 @@ class AnthropicAdapter(LLMInterface):
name = "Anthropic"
model: str
default_instructor_mode = "anthropic_tools"
def __init__(self, max_completion_tokens: int, model: str = None):
import anthropic
config_instructor_mode = get_llm_config().llm_instructor_mode
instructor_mode = (
config_instructor_mode if config_instructor_mode else self.default_instructor_mode
)
self.aclient = instructor.patch(
create=anthropic.AsyncAnthropic(api_key=get_llm_config().llm_api_key).messages.create,
mode=instructor.Mode.ANTHROPIC_TOOLS,
mode=instructor.Mode(instructor_mode),
)
self.model = model

12
cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py
View file

@ -41,6 +41,7 @@ class GeminiAdapter(LLMInterface):
name: str
model: str
api_key: str
default_instructor_mode = "json_mode"
def __init__(
self,
@ -63,7 +64,16 @@ class GeminiAdapter(LLMInterface):
self.fallback_api_key = fallback_api_key
self.fallback_endpoint = fallback_endpoint
self.aclient = instructor.from_litellm(litellm.acompletion, mode=instructor.Mode.JSON)
from cognee.infrastructure.llm.config import get_llm_config
config_instructor_mode = get_llm_config().llm_instructor_mode
instructor_mode = (
config_instructor_mode if config_instructor_mode else self.default_instructor_mode
)
self.aclient = instructor.from_litellm(
litellm.acompletion, mode=instructor.Mode(instructor_mode)
)
@retry(
stop=stop_after_delay(128),

12
cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py
View file

@ -41,6 +41,7 @@ class GenericAPIAdapter(LLMInterface):
name: str
model: str
api_key: str
default_instructor_mode = "json_mode"
def __init__(
self,
@ -63,7 +64,16 @@ class GenericAPIAdapter(LLMInterface):
self.fallback_api_key = fallback_api_key
self.fallback_endpoint = fallback_endpoint
self.aclient = instructor.from_litellm(litellm.acompletion, mode=instructor.Mode.JSON)
from cognee.infrastructure.llm.config import get_llm_config
config_instructor_mode = get_llm_config().llm_instructor_mode
instructor_mode = (
config_instructor_mode if config_instructor_mode else self.default_instructor_mode
)
self.aclient = instructor.from_litellm(
litellm.acompletion, mode=instructor.Mode(instructor_mode)
)
@retry(
stop=stop_after_delay(128),

8
cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/mistral/adapter.py
View file

@ -37,6 +37,7 @@ class MistralAdapter(LLMInterface):
model: str
api_key: str
max_completion_tokens: int
default_instructor_mode = "mistral_tools"
def __init__(self, api_key: str, model: str, max_completion_tokens: int, endpoint: str = None):
from mistralai import Mistral
@ -44,9 +45,14 @@ class MistralAdapter(LLMInterface):
self.model = model
self.max_completion_tokens = max_completion_tokens
config_instructor_mode = get_llm_config().llm_instructor_mode
instructor_mode = (
config_instructor_mode if config_instructor_mode else self.default_instructor_mode
)
self.aclient = instructor.from_litellm(
litellm.acompletion,
mode=instructor.Mode.MISTRAL_TOOLS,
mode=instructor.Mode(instructor_mode),
api_key=get_llm_config().llm_api_key,
)

12
cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py
View file

@ -42,6 +42,8 @@ class OllamaAPIAdapter(LLMInterface):
- aclient
"""
default_instructor_mode = "json_mode"
def __init__(
self, endpoint: str, api_key: str, model: str, name: str, max_completion_tokens: int
):
@ -51,8 +53,16 @@ class OllamaAPIAdapter(LLMInterface):
self.endpoint = endpoint
self.max_completion_tokens = max_completion_tokens
from cognee.infrastructure.llm.config import get_llm_config
config_instructor_mode = get_llm_config().llm_instructor_mode
instructor_mode = (
config_instructor_mode if config_instructor_mode else self.default_instructor_mode
)
self.aclient = instructor.from_openai(
OpenAI(base_url=self.endpoint, api_key=self.api_key), mode=instructor.Mode.JSON
OpenAI(base_url=self.endpoint, api_key=self.api_key),
mode=instructor.Mode(instructor_mode),
)
@retry(

11
cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py
View file

@ -56,6 +56,7 @@ class OpenAIAdapter(LLMInterface):
model: str
api_key: str
api_version: str
default_instructor_mode = "json_schema_mode"
MAX_RETRIES = 5
@ -74,14 +75,20 @@ class OpenAIAdapter(LLMInterface):
fallback_api_key: str = None,
fallback_endpoint: str = None,
):
from cognee.infrastructure.llm.config import get_llm_config
config_instructor_mode = get_llm_config().llm_instructor_mode
instructor_mode = (
config_instructor_mode if config_instructor_mode else self.default_instructor_mode
)
# TODO: With gpt5 series models OpenAI expects JSON_SCHEMA as a mode for structured outputs.
# Make sure all new gpt models will work with this mode as well.
if "gpt-5" in model:
self.aclient = instructor.from_litellm(
litellm.acompletion, mode=instructor.Mode.JSON_SCHEMA
litellm.acompletion, mode=instructor.Mode(instructor_mode)
)
self.client = instructor.from_litellm(
litellm.completion, mode=instructor.Mode.JSON_SCHEMA
litellm.completion, mode=instructor.Mode(instructor_mode)
)
else:
self.aclient = instructor.from_litellm(litellm.acompletion)