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

Update adapter.py

Browse source
This commit is contained in:
Vasilije 2025-01-16 13:26:37 +01:00 • committed by GitHub
parent b6e82dfb4f
commit 6c1c8abc26
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 12 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
cognee/infrastructure/llm/openai/adapter.py
View file

@ -12,7 +12,6 @@ from cognee.infrastructure.llm.llm_interface import LLMInterface
from cognee.infrastructure.llm.prompts import read_query_prompt from cognee.infrastructure.llm.prompts import read_query_prompt
from cognee.base_config import get_base_config from cognee.base_config import get_base_config
monitoring = get_base_config().monitoring_tool monitoring = get_base_config().monitoring_tool
if monitoring == MonitoringTool.LANGFUSE: if monitoring == MonitoringTool.LANGFUSE:
from langfuse.decorators import observe from langfuse.decorators import observe
@ -35,14 +34,19 @@ class OpenAIAdapter(LLMInterface):
transcription_model: str, transcription_model: str,
streaming: bool = False, streaming: bool = False,
): ):
self.aclient = instructor.from_litellm(litellm.acompletion)
self.client = instructor.from_litellm(litellm.completion)
self.transcription_model = transcription_model
self.model = model
self.api_key = api_key
self.endpoint = endpoint
self.api_version = api_version
self.streaming = streaming
@observe(as_type="generation")
async def acreate_structured_output(
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
@observe(as_type='generation') ) -> BaseModel:
async def acreate_structured_output(self, text_input: str, system_prompt: str,
response_model: Type[BaseModel]) -> BaseModel:
"""Generate a response from a user query.""" """Generate a response from a user query."""
return await self.aclient.chat.completions.create( return await self.aclient.chat.completions.create(