dnd-helpers/src/llm/processor.py

import os
from typing import Any, Dict, Optional

from openai import OpenAI
from pydantic import ValidationError

from .models import ExtractionResult
from .prompts import EXTRACTION_SYSTEM_PROMPT, NOISE_FILTER_SYSTEM_PROMPT


class LLMProcessor:
    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: Optional[str] = None,
    ):
        """
        Initializes the LLMProcessor.

        :param api_key: OpenAI API key. If None, it looks for OPENAI_API_KEY in environment variables.
        :param base_url: OpenAI-compatible base URL (e.g., for vLLM).
        :param model: The model to use for processing. If None, it looks for LLM_MODEL in environment variables.
        """
        self.client = OpenAI(
            api_key=api_key or os.environ.get("OPENAI_API_KEY"),
            base_url=base_url or os.environ.get("OPENAI_BASE_URL"),
        )
        self.model = model or os.environ.get("LLM_MODEL", "gpt-4o")

    def _call_llm(
        self,
        system_prompt: str,
        user_prompt: str,
        response_format: Optional[Any] = None,
    ) -> str:
        """
        Generic method to call the LLM.
        """
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                response_format=response_format,
                extra_body={"include_reasoning": False},
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"LLM Error: {e}")
            return ""

    def filter_transcript(self, text: str) -> str:
        """
        Stage 1: Raw Transcript -> Filtered Text.
        """
        result = self._call_llm(NOISE_FILTER_SYSTEM_PROMPT, text)
        print(f"LLM Processor (Filter): {text} -> {result}")
        return result

    def extract_structured_data(self, filtered_text: str) -> ExtractionResult:
        """
        Stage 2: Filtered Text -> Structured Data.
        """
        print(f"LLM Processor (Extract): Calling extraction for: {filtered_text}")
        try:
            # Using standard chat.completions.create with JSON mode for better compatibility with vLLM
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": EXTRACTION_SYSTEM_PROMPT},
                    {"role": "user", "content": filtered_text},
                ],
                response_format={"type": "json_object"},
                extra_body={"include_reasoning": False},
            )

            import json

            content = response.choices[0].message.content
            print(f"LLM Processor (Extract): Raw JSON response: {content}")
            data = json.loads(content)

            # Map the JSON data to the Pydantic model
            return ExtractionResult(**data)

        except Exception as e:
            print(f"Extraction Error: {e}")
            # Return an empty ExtractionResult if parsing fails
            return ExtractionResult()

    def process_pipeline(self, raw_text: str) -> ExtractionResult:
        """
        Executes the two-stage pipeline: Raw Transcript -> Filtered Text -> Structured Data.
        """
        filtered_text = self.filter_transcript(raw_text)
        if not filtered_text:
            return ExtractionResult()

        return self.extract_structured_data(filtered_text)