Fix langfuse

This commit is contained in:
vasilije 2024-12-18 19:01:29 +01:00
parent 92ecd8a024
commit c448dfb96d
5 changed files with 509 additions and 575 deletions

View file

@ -17,6 +17,10 @@ from cognee.modules.pipelines.operations.log_pipeline_status import log_pipeline
from cognee.tasks.documents import classify_documents, check_permissions_on_documents, extract_chunks_from_documents from cognee.tasks.documents import classify_documents, check_permissions_on_documents, extract_chunks_from_documents
from cognee.tasks.graph import extract_graph_from_code from cognee.tasks.graph import extract_graph_from_code
from cognee.tasks.storage import add_data_points from cognee.tasks.storage import add_data_points
from cognee.base_config import get_base_config
from cognee.shared.data_models import MonitoringTool
if MonitoringTool.LANGFUSE:
from langfuse.decorators import observe
logger = logging.getLogger("code_graph_pipeline") logger = logging.getLogger("code_graph_pipeline")
@ -49,7 +53,7 @@ async def code_graph_pipeline(datasets: Union[str, list[str]] = None, user: User
return await asyncio.gather(*awaitables) return await asyncio.gather(*awaitables)
@observe
async def run_pipeline(dataset: Dataset, user: User): async def run_pipeline(dataset: Dataset, user: User):
data_documents: list[Data] = await get_dataset_data(dataset_id = dataset.id) data_documents: list[Data] = await get_dataset_data(dataset_id = dataset.id)

View file

@ -10,7 +10,9 @@ class BaseConfig(BaseSettings):
monitoring_tool: object = MonitoringTool.LANGFUSE monitoring_tool: object = MonitoringTool.LANGFUSE
graphistry_username: Optional[str] = os.getenv("GRAPHISTRY_USERNAME") graphistry_username: Optional[str] = os.getenv("GRAPHISTRY_USERNAME")
graphistry_password: Optional[str] = os.getenv("GRAPHISTRY_PASSWORD") graphistry_password: Optional[str] = os.getenv("GRAPHISTRY_PASSWORD")
langfuse_public_key: Optional[str] = os.getenv("LANGFUSE_PUBLIC_KEY")
langfuse_secret_key: Optional[str] = os.getenv("LANGFUSE_SECRET_KEY")
langfuse_host: Optional[str] = os.getenv("LANGFUSE_HOST")
model_config = SettingsConfigDict(env_file = ".env", extra = "allow") model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
def to_dict(self) -> dict: def to_dict(self) -> dict:

View file

@ -6,10 +6,14 @@ from typing import Type
import litellm import litellm
import instructor import instructor
from pydantic import BaseModel from pydantic import BaseModel
from cognee.shared.data_models import MonitoringTool
from cognee.exceptions import InvalidValueError from cognee.exceptions import InvalidValueError
from cognee.infrastructure.llm.llm_interface import LLMInterface 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
if MonitoringTool.LANGFUSE:
from langfuse.decorators import observe
class OpenAIAdapter(LLMInterface): class OpenAIAdapter(LLMInterface):
name = "OpenAI" name = "OpenAI"
@ -18,6 +22,7 @@ class OpenAIAdapter(LLMInterface):
api_version: str api_version: str
"""Adapter for OpenAI's GPT-3, GPT=4 API""" """Adapter for OpenAI's GPT-3, GPT=4 API"""
def __init__( def __init__(
self, self,
api_key: str, api_key: str,
@ -35,8 +40,13 @@ class OpenAIAdapter(LLMInterface):
self.endpoint = endpoint self.endpoint = endpoint
self.api_version = api_version self.api_version = api_version
self.streaming = streaming self.streaming = streaming
base_config = get_base_config()
@observe()
async def acreate_structured_output(self, text_input: str, system_prompt: str,
response_model: Type[BaseModel]) -> 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(
@ -56,7 +66,9 @@ class OpenAIAdapter(LLMInterface):
max_retries=5, max_retries=5,
) )
def create_structured_output(self, text_input: str, system_prompt: str, response_model: Type[BaseModel]) -> BaseModel: @observe
def create_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 self.client.chat.completions.create( return self.client.chat.completions.create(

993
poetry.lock generated

File diff suppressed because it is too large Load diff

View file

@ -46,7 +46,7 @@ aiofiles = "^23.2.1"
qdrant-client = {version = "^1.9.0", optional = true} qdrant-client = {version = "^1.9.0", optional = true}
graphistry = "^0.33.5" graphistry = "^0.33.5"
tenacity = "^8.4.1" tenacity = "^8.4.1"
weaviate-client = {version = "4.6.7", optional = true} weaviate-client = {version = "4.9.6", optional = true}
scikit-learn = "^1.5.0" scikit-learn = "^1.5.0"
pypdf = "^4.1.0" pypdf = "^4.1.0"
neo4j = {version = "^5.20.0", optional = true} neo4j = {version = "^5.20.0", optional = true}
@ -74,6 +74,7 @@ deepeval = {version = "^2.0.1", optional = true}
transformers = "^4.46.3" transformers = "^4.46.3"
pymilvus = {version = "^2.5.0", optional = true} pymilvus = {version = "^2.5.0", optional = true}
unstructured = { extras = ["csv", "doc", "docx", "epub", "md", "odt", "org", "ppt", "pptx", "rst", "rtf", "tsv", "xlsx"], version = "^0.16.10", optional = true } unstructured = { extras = ["csv", "doc", "docx", "epub", "md", "odt", "org", "ppt", "pptx", "rst", "rtf", "tsv", "xlsx"], version = "^0.16.10", optional = true }
httpx = "0.27.0"