dv-llmclient 0.1.1__tar.gz

dv_llmclient-0.1.1/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.3
+Name: dv-llmclient
+Version: 0.1.1
+Summary: An LLM client with cache and type checking
+Author: pnmougel
+Author-email: pnmougel <pierre-nicolas.mougel@datavalor.com>
+Requires-Dist: hishel>=1.3.0
+Requires-Dist: langfuse>=4.14.1
+Requires-Dist: pydantic>=2.13.4
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+

dv_llmclient-0.1.1/pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "dv-llmclient"
+version = "0.1.1"
+description = "An LLM client with cache and type checking"
+readme = "README.md"
+authors = [
+    { name = "pnmougel", email = "pierre-nicolas.mougel@datavalor.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "hishel>=1.3.0",
+    "langfuse>=4.14.1",
+    "pydantic>=2.13.4",
+]
+
+[build-system]
+requires = ["uv_build>=0.11.21,<0.12.0"]
+build-backend = "uv_build"

dv_llmclient-0.1.1/src/dv_llmclient/__init__.py

@@ -0,0 +1,4 @@
+from .models import Endpoint, ReasoningEffort, LlmModel, ListModel
+from .prompt import Prompt, PromptPart, FilePart, BinaryPart, ImagePart, LfPrompt
+from .response_validator import CompletionResponseValidator, ResponseValidator
+from .llm_client import LlmClient

dv_llmclient-0.1.1/src/dv_llmclient/exceptions.py

@@ -0,0 +1,18 @@
+from typing import Optional
+
+class LLMClientError(Exception):
+    """Base exception class for LLM client errors."""
+    def __init__(self, message: str, base_exception: Optional[Exception] = None):
+        super().__init__(message)
+        self.message = message
+        self.base_exception = base_exception
+
+class LLMClientSchemaBuildError(LLMClientError):
+    def __init__(self, base_exception: Exception, model):
+        super().__init__(f"Error building schema for model: {model}", base_exception)
+        self.model = model
+
+class LLMClientMaxRetryError(LLMClientError):
+    def __init__(self, max_retries: int):
+        super().__init__(f"Maximum retries reached ({max_retries})")
+

dv_llmclient-0.1.1/src/dv_llmclient/lf_client.py

@@ -0,0 +1,4 @@
+from langfuse import Langfuse
+import os
+
+lf_client = Langfuse(public_key=os.getenv("LANGFUSE_PUBLIC_KEY"), secret_key=os.getenv("LANGFUSE_SECRET_KEY"), media_upload_thread_count=5)

dv_llmclient-0.1.1/src/dv_llmclient/llm_client.py

@@ -0,0 +1,213 @@
+from pathlib import Path
+
+from langfuse import LangfuseGeneration
+from pydantic import BaseModel
+from typing import TypeVar, Generic, Type, Optional, Union
+import hishel
+import json
+import os
+import base64
+from hishel.httpx import SyncCacheClient
+
+from .lf_client import lf_client
+from .exceptions import LLMClientMaxRetryError, LLMClientSchemaBuildError
+from .models import ImageOCR, OCRResponse, LlmModel, ReasoningEffort, Endpoint
+from .prompt import ImagePart, LfPrompt, Prompt, PromptPart
+from .response_validator import OcrResponseValidator, RequestCacheFilter, ResponseValidator, CompletionResponseValidator
+
+M = TypeVar("M", bound=Union[str, BaseModel])
+
+class LlmClient(Generic[M]):
+    def __init__(
+        self,
+        model: LlmModel,
+        output_type: Type[M] = str,
+        temperature: Optional[float] = None,
+        max_retries: int = 3,
+        reasoning: Optional[ReasoningEffort] = None,
+        endpoint: Endpoint = "eden",
+        extra_body: dict = {},
+        validation: Optional[ResponseValidator[M]] = None,
+        default_on_error: Optional[M] = None,
+        force_cache_refresh: bool = False,
+        disable_cache: bool = False
+    ):
+        self.model = model
+        self.output_type = output_type
+        self.temperature = temperature
+        self.max_retries = max_retries
+        self.reasoning = reasoning
+        self.endpoint = endpoint
+        self.extra_body = extra_body
+        self.validation = validation
+        self.default_on_error = default_on_error
+        self.force_cache_refresh = force_cache_refresh
+        self.disable_cache = disable_cache
+
+    def ocr(self, pdf_path: str | Path, extract_header: bool = True, extract_footer: bool = True, table_format: str = "markdown", include_image_base64: bool = True) -> M:
+        with lf_client.start_as_current_generation(name="llm_call", model=self.model, 
+                                                    metadata={"endpoint": self.endpoint, "max_retries": self.max_retries, "cache_disabled": self.disable_cache, "cache_refresh": self.force_cache_refresh},
+                                                    model_parameters={"temperature": str(self.temperature), "reasoning": self.reasoning}) as span:
+            document_url = f"data:application/pdf;base64,{base64.b64encode(Path(pdf_path).read_bytes()).decode()}"
+            bbox_annotation_format = {"type":"json_schema","json_schema":{"name":"ImageOCR","schema":ImageOCR.model_json_schema(),"strict":True}}
+            query = dict(
+                model=self.model,
+                extract_header=extract_header,
+                extract_footer=extract_footer,
+                table_format=table_format,
+                document={"type": "document_url", "document_url": document_url},
+                bbox_annotation_format=bbox_annotation_format, # type: ignore
+                include_image_base64=include_image_base64
+            )
+            return self.__run(query=query, span=span, path="/v1/ocr", validator=OcrResponseValidator()) # type: ignore
+
+    def completion(self, prompt: Prompt | PromptPart | list[PromptPart]) -> M:
+        with lf_client.start_as_current_generation(name="llm_call", model=self.model, 
+                                                    metadata={"endpoint": self.endpoint, "max_retries": self.max_retries, "cache_disabled": self.disable_cache, "cache_refresh": self.force_cache_refresh},
+                                                    model_parameters={"temperature": str(self.temperature), "reasoning": self.reasoning}) as span:
+            response_format: dict[str, str] = {"type": "text"}
+            try:
+                if self.output_type != str:
+                    response_format = { "type": "json_schema", "json_schema": self.output_type.model_json_schema()} # type: ignore 
+            except Exception as e:
+                # lf_client.create_event(name="output_schema_error", status_message="Error while building output schema", level="ERROR", metadata={"error": str(e)})
+                raise LLMClientSchemaBuildError(e, self.output_type)
+            span.update_trace(metadata={"response_format_type": response_format['type']})
+            span.update_trace(metadata={"response_format": response_format})
+
+            if not isinstance(prompt, Prompt):
+                if not isinstance(prompt, list):
+                    prompt = [prompt]
+                prompt = Prompt(*prompt)
+            for part in prompt.parts:
+                if isinstance(part, LfPrompt):
+                    span.update(prompt=part.lf_prompt)
+                if isinstance(part, ImagePart):
+                    lf_client.update_current_trace( metadata={"document": part.as_lf_media()} )
+                span.update(input=prompt.as_log_output())
+            
+            query = {
+                "model": self.model,
+                "temperature": self.temperature,
+                "response_format": response_format,
+                "messages": [prompt.build()],
+                "reasoning_effort": self.reasoning,
+                "temperature": self.temperature,
+            }
+            validator: ResponseValidator[M] = (
+                CompletionResponseValidator(output_type=self.output_type)
+                if self.validation is None
+                else self.validation
+            )
+            path = "/v3/llm/chat/completions" if self.endpoint == 'eden' else "/v1/chat/completions"
+
+            return self.__run(query=query, span=span, path=path, validator=validator)
+
+    def __run(
+        self,
+        query: dict[str, Union[str, bool, dict]],
+        path: str,
+        span: LangfuseGeneration,
+        validator: ResponseValidator[M]
+    ) -> M:
+        try:
+            self.cache_storage = hishel.SyncSqliteStorage(database_path="./cache.db", default_ttl=3600*24*30)
+
+            # Build the response validator for both the http cache
+            # validator: ResponseValidator[M] = (
+            #     CompletionResponseValidator(output_type=self.output_type)
+            #     if self.validation is None
+            #     else self.validation
+            # )
+            request_cache_filter = RequestCacheFilter(cache_storage=self.cache_storage, force_refresh=self.force_cache_refresh, disable_cache=self.disable_cache)
+
+            base_url_mapping = {
+                "openrouter": "https://openrouter.ai/api",
+                "dgx": "http://localhost:12000/api",
+                "test": "http://localhost:8000",
+                "mistral": "https://api.mistral.ai",
+                "eden": "https://api.edenai.run"
+            }
+            base_url = base_url_mapping.get(self.endpoint, self.endpoint)
+            api_key_mapping = {"openrouter": "OPENROUTER_API_KEY", "dgx": "OPEN_WEBUI_KEY", "mistral": "MISTRAL_API_KEY", "eden": "EDEN_AI_API_KEY"}
+            http_client = SyncCacheClient(
+                storage=self.cache_storage,
+                policy=hishel.FilterPolicy(request_filters=[request_cache_filter], response_filters=[validator])
+            )
+
+            nb_retry = 0
+            is_ok = False
+            res = None
+            while not is_ok and nb_retry < self.max_retries:
+                is_ok = True
+                with lf_client.start_as_current_observation(name="llm_query", as_type="generation", input=query) as llm_query_span:
+                    try:
+                        res = http_client.post(
+                            url=f"{base_url}{path}",
+                            content=json.dumps(query),
+                            headers={
+                                "Authorization": f"Bearer {os.getenv(api_key_mapping.get(self.endpoint, 'OPENAI_KEY'))}"
+                            },
+                            extensions={"hishel_body_key": True},
+                        )
+                        res = validator.build_response(res, None, llm_query_span)
+                    except Exception as e:
+                        print(f"Error while making request: {e}")
+                        llm_query_span.update(level="ERROR", status_message=f"Error during HTTP request: {e}")
+                        raise e
+                    if res is None:
+                        # Subsequent requests should not be loaded from cache, as the cache might be storing invalid responses
+                        # This is the case when the validation change for a given request that was previously cached. 
+                        # We must be able to reload a fresh response.
+                        request_cache_filter.force_refresh = True
+                        is_ok = False
+                        nb_retry += 1
+                
+            if nb_retry == self.max_retries:
+                if self.default_on_error is not None:
+                    return self.default_on_error
+                else:
+                    raise LLMClientMaxRetryError(self.max_retries)
+            
+            # Update the span with the clean output
+            if res is not None:
+                if self.output_type == str:
+                    span.update(output=res)
+                else:
+                    span.update(output=res.model_dump_json()) # type: ignore
+            return res # type: ignore
+        except Exception as e:
+            span.update(level="ERROR", status_message=f"{e}")
+            raise e
+        
+class LlmOCRClient(LlmClient[OCRResponse]):
+    def __init__(
+        self,
+        model: LlmModel = 'mistral-ocr-latest',
+        temperature: Optional[float] = None,
+        max_retries: int = 3,
+        reasoning: Optional[ReasoningEffort] = None,
+        endpoint: Endpoint = "mistral",
+        validation: Optional[ResponseValidator[OCRResponse]] = None,
+        default_on_error: Optional[OCRResponse] = None,
+        force_cache_refresh: bool = False,
+        disable_cache: bool = False
+    ):
+        super().__init__(
+            model=model,
+            output_type=OCRResponse,
+            temperature=temperature,
+            max_retries=max_retries,
+            reasoning=reasoning,
+            endpoint=endpoint,
+            validation=validation,
+            default_on_error=default_on_error,
+            force_cache_refresh=force_cache_refresh,
+            disable_cache=disable_cache
+        )
+
+    def ocr(self, pdf_path: str | Path, extract_header: bool = True, extract_footer: bool = True, table_format: str = "markdown", include_image_base64: bool = True) -> OCRResponse:
+        return super().ocr(pdf_path=pdf_path, extract_header=extract_header, extract_footer=extract_footer, table_format=table_format, include_image_base64=include_image_base64)
+    
+    def completion(self, prompt):
+        raise NotImplementedError("completion is not implemented for LlmOCR, use the ocr method instead.")

dv_llmclient-0.1.1/src/dv_llmclient/models.py

@@ -0,0 +1,127 @@
+from enum import Enum
+
+from pydantic import BaseModel, Field, RootModel
+
+from pydantic import BaseModel
+from typing import Iterator, TypeVar, List, Optional, Literal
+
+type ReasoningEffort = Literal["none", "low", "medium", "high"]
+type Endpoint = Literal["dgx", "openrouter", "test", "mistral", "eden"] | str
+type LlmModel = Literal[
+    "google/gemini-2.5-flash",
+    "google/gemini-2.5-pro",
+    "google/gemini-3-pro-preview",
+    "gpt-oss:120b",
+    "mistral-ocr-latest",
+] | str
+
+T = TypeVar("T", bound=BaseModel)
+
+class ListModel[T](RootModel[List[T]]):
+    """
+    Helper class to handle responses that are lists of items. It provides list-like access to the items while still being a Pydantic model.
+    """
+    root: list[T]
+
+    def __iter__(self) -> Iterator[T]: # type: ignore
+        return iter(self.root)
+
+    def __getitem__(self, item: int) -> T:
+        return self.root[item]
+
+    def __len__(self) -> int:
+        return len(self.root)
+    
+    def aslist(self) -> list[T]:
+        return self.root
+
+
+class LLMResponseMessage(BaseModel):
+    role: str
+    content: str
+    reasoning_context: Optional[str] = None
+    reasoning: Optional[str] = None
+
+
+class LLMResponseChoice(BaseModel):
+    finish_reason: str
+    message: LLMResponseMessage
+
+
+class LLMResponseCompletionTokensDetails(BaseModel):
+    reasoning_tokens: int
+    image_tokens: Optional[int] = None
+
+
+class LLMResponseCostDetails(BaseModel):
+    upstream_inference_cost: Optional[float] = None
+    upstream_inference_prompt_cost: Optional[float] = None
+    upstream_inference_completions_cost: Optional[float] = None  
+
+class LLMResponseUsage(BaseModel):
+    cost: Optional[float] = None
+    completion_tokens: int
+    prompt_tokens: int
+    total_tokens: int
+    completion_tokens_details: LLMResponseCompletionTokensDetails
+    cost_details: Optional[LLMResponseCostDetails] = None
+
+
+class LLMResponseRaw(BaseModel):
+    choices: List[LLMResponseChoice]
+    usage: LLMResponseUsage
+    cost: Optional[float] = None
+
+
+# OCR related models
+class ImageType(str, Enum):
+    GRAPH = "graph"
+    TEXT = "text"
+    TABLE = "table"
+    IMAGE = "image"
+
+class ImageOCR(BaseModel):
+    image_type: ImageType = Field(..., description="Type de l'image. Doit être 'graph', 'text', 'table' ou 'image'.")
+    caption: str = Field(..., description="Légende de la figure, si disponible.")
+    description: str = Field(..., description="Description de l'image.")
+
+class OCRUsageInfo(BaseModel):
+    pages_processed: int
+    doc_size_bytes: int
+
+class OCRTable(BaseModel):
+    id: str
+    content: str
+    format: str
+    # word_confidence_score: Optional[float]
+
+class OCRPageDimension(BaseModel):
+    dpi: int
+    height: int
+    width: int
+
+
+class OCRImageResponse(BaseModel):
+    id: str
+    bottom_right_x: Optional[int]
+    bottom_right_y: Optional[int]
+    top_left_x: Optional[int]
+    top_left_y: Optional[int]
+    image_annotation: Optional[str]
+    image_base64: Optional[str]
+
+class OCRPage(BaseModel):
+    index: int
+    markdown: str
+    images: list[OCRImageResponse]
+    tables: list[OCRTable]
+    header: Optional[str]
+    footer: Optional[str]
+    hyperlinks: Optional[list[str]]
+    dimensions: Optional[OCRPageDimension]
+    confidence_scores: Optional[float]
+
+class OCRResponse(BaseModel):
+    usage_info: Optional[OCRUsageInfo]
+    pages: list[OCRPage]
+    document_annotation: Optional[str]

dv_llmclient-0.1.1/src/dv_llmclient/prompt.py

@@ -0,0 +1,100 @@
+import base64
+from pathlib import Path
+from typing import Optional
+from mimetypes import guess_type
+from langfuse import get_client
+from langfuse.media import LangfuseMedia
+from lf.client import lf_client
+
+class FilePart:
+    def __init__(self, file_path: str | Path):
+        self.file_path = Path(file_path) if isinstance(file_path, str) else file_path
+        self.media_type = guess_type(self.file_path)[0]
+        self.base64_content = base64.b64encode(self.file_path.read_bytes()).decode()
+
+    def as_dict(self) -> dict:
+        return {
+            "file": {
+                "file_data": f"data:{self.media_type};base64,{self.base64_content}",
+                "filename": self.file_path.name,
+            },
+            "type": "file",
+        }
+    
+    def __repr__(self):
+        return f"<FilePart filename={self.file_path.name} media_type={self.media_type} b64_size={len(self.base64_content)}>"
+
+class BinaryPart():
+    def __init__(self, data: bytes, filename: str, media_type: Optional[str] = None):
+        self.data = data
+        self.filename = filename
+        self.media_type = media_type or "application/octet-stream"
+        
+    def as_dict(self) -> dict:
+        return {
+            "file": {
+                "file_data": f"data:{self.media_type};base64,{base64.b64encode(self.data).decode()}",
+                "filename": self.filename,
+            },
+            "type": "file",
+        }
+    
+    def __repr__(self):
+        return f"<BinaryPart filename={self.filename} media_type={self.media_type}>"
+
+class ImagePart():
+    def __init__(self, file_path: Path, media_type: Optional[str] = None):
+        self.file_path = file_path
+        self.media_type = media_type or guess_type(file_path)[0]
+        
+    def as_dict(self) -> dict:
+        base64_content = base64.b64encode(self.file_path.read_bytes()).decode()
+        return {
+            # "image_url": {
+            #     "url": f"data:{self.media_type};base64,{base64_content}",
+            # },
+            "image_url": f"data:{self.media_type};base64,{base64_content}",
+            "type": "image_url",
+        }
+    
+    def as_lf_media(self) -> LangfuseMedia:
+        return LangfuseMedia(
+            content_type=self.media_type,  # type: ignore
+            content_bytes=self.file_path.read_bytes(),
+            file_path=self.file_path.name
+        )
+
+    def __repr__(self):
+        return f"<ImagePart media_type={self.media_type} path={self.file_path}>"
+
+class LfPrompt:
+    def __init__(self, name: str, label: Optional[str] = "latest", params: dict = {}):
+        self.lf_prompt = lf_client.get_prompt(name, label=label, cache_ttl_seconds=0)
+        self.text = self.lf_prompt.compile(**params)
+
+    def as_dict(self) -> dict:
+        return {"type": "text", "text": self.text}
+
+    def __repr__(self):
+        return self.text
+
+
+type PromptPart = FilePart | BinaryPart | ImagePart | LfPrompt | str
+
+class Prompt():
+    def __init__(self, *args: PromptPart):
+        self.parts: list[FilePart | BinaryPart | ImagePart | LfPrompt | str] = list(args)
+
+    def build(self) -> dict:
+        content = [{"type": "text", "text": part} if isinstance(part, str) else part.as_dict() for part in self.parts]
+        return {"role": "user", "content": content}
+    
+    def add_part(self, part: PromptPart):
+        self.parts.append(part)
+
+    def __repr__(self):
+        prefix = f"<Prompt {len(self.parts)} parts>"
+        return prefix + "\n" + "\n".join([f" - {repr(part)}" for part in self.parts])
+    
+    def as_log_output(self) -> str:
+        return "\n".join([repr(part) for part in self.parts])

dv_llmclient-0.1.1/src/dv_llmclient/response_validator.py

@@ -0,0 +1,186 @@
+
+from typing import Iterable, Iterator, Optional, Type
+
+import hishel
+import httpx
+import gzip
+import hashlib
+
+from langfuse import LangfuseGeneration, LangfuseSpan
+from pydantic_core import ValidationError
+
+from .models import LLMResponseRaw, OCRResponse
+from hishel._utils import make_sync_iterator
+from abc import ABC, abstractmethod
+
+class RequestCacheFilter(hishel.BaseFilter[hishel.Request]):
+    def __init__(self, cache_storage: hishel.SyncBaseStorage, force_refresh: bool = True, disable_cache: bool = False):
+        self.cache_storage = cache_storage
+        self.force_refresh = force_refresh
+        self.disable_cache = disable_cache
+
+    def needs_body(self) -> bool:
+        return True
+
+    def apply(self, item: hishel.Request, body: bytes | None):
+        if self.disable_cache:
+            # Returning False completely disables the cache
+            return False
+        # If we just return False when force_refresh is False, the request will be sent but it wont be stored in cache...
+        # Instead we remove the cached response for the storage and always return True.
+        
+        # Dangerous : the key calculation is extracted from the hishel codebase, and might change in future versions.
+        if self.force_refresh:
+            assert isinstance(item.stream, (Iterator, Iterable))
+            collected = b"".join([chunk for chunk in item.stream])
+            key = hashlib.sha256(collected).hexdigest()
+            item.stream = make_sync_iterator([collected])
+            for entry in self.cache_storage.get_entries(key):
+                self.cache_storage.remove_entry(entry.id)
+        return True
+
+class ResponseValidator[M](hishel.BaseFilter[hishel.Response]):
+    def __init__(self, output_type: Type[M]):
+        self.output_type = output_type
+
+    def needs_body(self) -> bool:
+        return True
+
+    def apply(self, item: hishel.Response, body: bytes | None):
+        try:
+            res = self.build_response(item, body)
+            return res is not None
+        except Exception as e:
+            return False
+
+    def build_response(
+        self, response: hishel.Response | httpx.Response, body: Optional[bytes] = None, langfuse_span: Optional[LangfuseGeneration] = None
+    ) -> Optional[M]:
+        """
+        Build the response object from the http response. 
+
+        There is three possible outputs :
+        - None : the response is invalid, but it can be retried. This is the case when the response content from the LLM is not valid according to the validation rules.
+        - M : the response is valid, and can be returned as is.
+        - Exception : the response is invalid, and should not be retried.
+        """
+        is_post_validation = isinstance(response, httpx.Response)
+        if response.status_code >= 300:
+            # Http status 
+            if is_post_validation:
+                raise Exception(f"Invalid response status: {response.status_code}, body: {response.text}")
+            else:
+                return None
+        if isinstance(response, httpx.Response):
+            body_content = response.text
+        elif isinstance(response, hishel.Response):
+            try:
+                if body is None:
+                    return None
+                if response.headers.get("Content-Encoding", "") == "gzip":
+                    body = gzip.decompress(body)
+                body_content: str = body.decode("utf-8")
+            except Exception as e:
+                raise Exception(f"Unable to decode response body") from e
+
+        if langfuse_span is not None:
+            langfuse_span.update(output=body_content)
+
+        return self.validate_body_content(body_content, langfuse_span)
+        # # Handle the case where the response is wrapped in a markdown code block
+        # try:
+        #     llm_response = LLMResponseRaw.model_validate_json(body_content)
+        # except ValidationError as e:
+        #     raise Exception(f"Http response body is not a valid LLM response: {body_content}") from e
+
+        # if langfuse_span is not None:
+        #     usage_details = {
+        #         "total": llm_response.usage.total_tokens, "input": llm_response.usage.prompt_tokens, "output": llm_response.usage.completion_tokens, "reasoning_tokens": llm_response.usage.completion_tokens_details.reasoning_tokens
+        #     }
+        #     cost_details = {
+        #         "total": llm_response.usage.cost or 0,
+        #         "upstream_inference_cost": llm_response.usage.cost_details.upstream_inference_cost if llm_response.usage.cost_details else None,
+        #         "input": llm_response.usage.cost_details.upstream_inference_prompt_cost if llm_response.usage.cost_details else None,
+        #         "output": llm_response.usage.cost_details.upstream_inference_completions_cost if llm_response.usage.cost_details else None,
+        #     }
+        #     # cost_details = { "total": llm_response.usage.cost or 0 }
+        #     langfuse_span.update(usage_details=usage_details, cost_details=cost_details)
+
+        # try:
+        #     response_content = llm_response.choices[0].message.content
+        #     if self.output_type == str:
+        #         content_to_validate: M = str(response_content) # type: ignore
+        #     else:
+        #         start_idx = response_content.find("```json")
+        #         end_idx = response_content.find("```", start_idx + 7)
+        #         if start_idx != -1 and end_idx != -1:
+        #             response_content = response_content[start_idx + 7:end_idx].strip()
+        #         content_to_validate: M = self.output_type.model_validate_json(response_content) # type: ignore
+
+        #     res = self.validate(content_to_validate)
+        # except ValidationError as e:
+        #     return None
+        
+        # return res
+
+    @abstractmethod
+    def validate_body_content(self, body_content: str, langfuse_span: Optional[LangfuseGeneration]) -> Optional[M]:
+        pass
+
+    def validate(self, content: M) -> M:
+        # Implement your validation logic here
+        return content
+
+
+class CompletionResponseValidator[M](ResponseValidator[M]):
+    def __init__(self, output_type: Type[M]):
+        super().__init__(output_type)
+
+    def validate_body_content(self, body_content: str, langfuse_span: Optional[LangfuseGeneration]) -> Optional[M]:
+        # Handle the case where the response is wrapped in a markdown code block
+        try:
+            llm_response = LLMResponseRaw.model_validate_json(body_content)
+        except ValidationError as e:
+            raise Exception(f"Http response body is not a valid LLM response: {body_content}") from e
+
+        if langfuse_span is not None:
+            usage_details = {
+                "total": llm_response.usage.total_tokens, "input": llm_response.usage.prompt_tokens, "output": llm_response.usage.completion_tokens, "reasoning_tokens": llm_response.usage.completion_tokens_details.reasoning_tokens
+            }
+            cost_details = {
+                "total": llm_response.usage.cost or 0,
+                "upstream_inference_cost": llm_response.usage.cost_details.upstream_inference_cost if llm_response.usage.cost_details else None,
+                "input": llm_response.usage.cost_details.upstream_inference_prompt_cost if llm_response.usage.cost_details else None,
+                "output": llm_response.usage.cost_details.upstream_inference_completions_cost if llm_response.usage.cost_details else None,
+            }
+            # cost_details = { "total": llm_response.usage.cost or 0 }
+            langfuse_span.update(usage_details=usage_details, cost_details=cost_details)
+
+        try:
+            response_content = llm_response.choices[0].message.content
+            if self.output_type == str:
+                content_to_validate: M = str(response_content) # type: ignore
+            else:
+                start_idx = response_content.find("```json")
+                end_idx = response_content.find("```", start_idx + 7)
+                if start_idx != -1 and end_idx != -1:
+                    response_content = response_content[start_idx + 7:end_idx].strip()
+                content_to_validate: M = self.output_type.model_validate_json(response_content) # type: ignore
+
+            res = self.validate(content_to_validate)
+        except ValidationError as e:
+            return None
+        
+        return res
+
+class OcrResponseValidator(ResponseValidator[OCRResponse]):
+    def __init__(self):
+        super().__init__(OCRResponse)
+
+    def validate_body_content(self, body_content: str, langfuse_span: Optional[LangfuseGeneration]) -> Optional[OCRResponse]:
+        try:
+            llm_response = OCRResponse.model_validate_json(body_content)
+            return llm_response
+        except ValidationError as e:
+            raise Exception(f"Http response body is not a valid OCR response: {body_content}") from e
+