Update adapter.py

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

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.base_config import get_base_config
monitoring = get_base_config().monitoring_tool
if monitoring == MonitoringTool.LANGFUSE:
from langfuse.decorators import observe
@ -35,14 +34,19 @@ class OpenAIAdapter(LLMInterface):
transcription_model: str,
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]) -> BaseModel:
@observe(as_type="generation")
async def acreate_structured_output(
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
) -> BaseModel:
"""Generate a response from a user query."""
return await self.aclient.chat.completions.create(