push

Author: hassiebp

langfuse/_client/client.py

@@ -2055,14 +2055,14 @@ def create_score(
             new_body = ScoreBody(
                 id=score_id,
                 session_id=session_id,
-                dataset_run_id=dataset_run_id,
-                trace_id=trace_id,
+                datasetRunId=dataset_run_id,
+                traceId=trace_id,
                 observation_id=observation_id,
                 name=name,
                 value=value,
                 dataType=data_type,  # type: ignore
                 comment=comment,
-                config_id=config_id,
+                configId=config_id,
                 environment=self._environment,
                 metadata=metadata,
             )

langfuse/_task_manager/score_ingestion_consumer.py

@@ -7,23 +7,19 @@
 from typing import Any, List, Optional
 
 import backoff
-
-from ..version import __version__ as langfuse_version
-
-try:
-    import pydantic.v1 as pydantic
-except ImportError:
-    import pydantic  # type: ignore
+from pydantic import BaseModel
 
 from langfuse._utils.parse_error import handle_exception
 from langfuse._utils.request import APIError, LangfuseClient
 from langfuse._utils.serializer import EventSerializer
 
+from ..version import __version__ as langfuse_version
+
 MAX_EVENT_SIZE_BYTES = int(os.environ.get("LANGFUSE_MAX_EVENT_SIZE_BYTES", 1_000_000))
 MAX_BATCH_SIZE_BYTES = int(os.environ.get("LANGFUSE_MAX_BATCH_SIZE_BYTES", 2_500_000))
 
 
-class ScoreIngestionMetadata(pydantic.BaseModel):
+class ScoreIngestionMetadata(BaseModel):
     batch_size: int
     sdk_name: str
     sdk_version: str
@@ -78,8 +74,8 @@ def _next(self) -> list:
                 )
 
                 # convert pydantic models to dicts
-                if "body" in event and isinstance(event["body"], pydantic.BaseModel):
-                    event["body"] = event["body"].dict(exclude_none=True)
+                if "body" in event and isinstance(event["body"], BaseModel):
+                    event["body"] = event["body"].model_dump(exclude_none=True)
 
                 item_size = self._get_item_size(event)
 
@@ -156,7 +152,7 @@ def _upload_batch(self, batch: List[Any]) -> None:
             sdk_name="python",
             sdk_version=langfuse_version,
             public_key=self._public_key,
-        ).dict()
+        ).model_dump()
 
         @backoff.on_exception(
             backoff.expo, Exception, max_tries=self._max_retries, logger=None

tests/utils.py

@@ -1,14 +1,8 @@
 import base64
 import os
-import typing
 from time import sleep
 from uuid import uuid4
 
-try:
-    import pydantic.v1 as pydantic  # type: ignore
-except ImportError:
-    import pydantic  # type: ignore
-
 from langfuse.api import LangfuseAPI
 
 
@@ -26,62 +20,6 @@ def get_api():
     )
 
 
-class LlmUsageWithCost(pydantic.BaseModel):
-    prompt_tokens: typing.Optional[int] = pydantic.Field(
-        alias="promptTokens", default=None
-    )
-    completion_tokens: typing.Optional[int] = pydantic.Field(
-        alias="completionTokens", default=None
-    )
-    total_tokens: typing.Optional[int] = pydantic.Field(
-        alias="totalTokens", default=None
-    )
-    input_cost: typing.Optional[float] = pydantic.Field(alias="inputCost", default=None)
-    output_cost: typing.Optional[float] = pydantic.Field(
-        alias="outputCost", default=None
-    )
-    total_cost: typing.Optional[float] = pydantic.Field(alias="totalCost", default=None)
-
-
-class CompletionUsage(pydantic.BaseModel):
-    completion_tokens: int
-    """Number of tokens in the generated completion."""
-
-    prompt_tokens: int
-    """Number of tokens in the prompt."""
-
-    total_tokens: int
-    """Total number of tokens used in the request (prompt + completion)."""
-
-
-class LlmUsage(pydantic.BaseModel):
-    prompt_tokens: typing.Optional[int] = pydantic.Field(
-        alias="promptTokens", default=None
-    )
-    completion_tokens: typing.Optional[int] = pydantic.Field(
-        alias="completionTokens", default=None
-    )
-    total_tokens: typing.Optional[int] = pydantic.Field(
-        alias="totalTokens", default=None
-    )
-
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {
-            "by_alias": True,
-            "exclude_unset": True,
-            **kwargs,
-        }
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {
-            "by_alias": True,
-            "exclude_unset": True,
-            **kwargs,
-        }
-        return super().dict(**kwargs_with_defaults)
-
-
 def encode_file_to_base64(image_path) -> str:
     with open(image_path, "rb") as file:
         return base64.b64encode(file.read()).decode("utf-8")