Temp
This commit is contained in:
Binary file not shown.
Binary file not shown.
+10
-3
@@ -46,11 +46,18 @@ class CharacterStateUpdate(BaseModel):
|
||||
|
||||
class ExtractionResult(BaseModel):
|
||||
lore_updates: List[LoreUpdate] = Field(
|
||||
default_factory=list, description="List of discovered lore facts"
|
||||
default_factory=list, description="List of discovered lore facts", alias="lore"
|
||||
)
|
||||
character_updates: List[CharacterStateUpdate] = Field(
|
||||
default_factory=list, description="List of character state changes"
|
||||
default_factory=list,
|
||||
description="List of character state changes",
|
||||
alias="character_state",
|
||||
)
|
||||
significant_events: List[str] = Field(
|
||||
default_factory=list, description="List of significant plot points or events"
|
||||
default_factory=list,
|
||||
description="List of significant plot points or events",
|
||||
alias="events",
|
||||
)
|
||||
|
||||
class Config:
|
||||
populate_by_name = True
|
||||
|
||||
+22
-12
@@ -13,20 +13,20 @@ class LLMProcessor:
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
model: str = "gpt-4o",
|
||||
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.
|
||||
: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
|
||||
self.model = model or os.environ.get("LLM_MODEL", "gpt-4o")
|
||||
|
||||
def _call_llm(
|
||||
self,
|
||||
@@ -45,6 +45,7 @@ class LLMProcessor:
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
response_format=response_format,
|
||||
extra_body={"include_reasoning": False},
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
except Exception as e:
|
||||
@@ -55,27 +56,36 @@ class LLMProcessor:
|
||||
"""
|
||||
Stage 1: Raw Transcript -> Filtered Text.
|
||||
"""
|
||||
return self._call_llm(NOISE_FILTER_SYSTEM_PROMPT, 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.
|
||||
"""
|
||||
# We use OpenAI's structured output (JSON mode/tool calling) via Pydantic's response_format.
|
||||
# For models that support it, we can pass the Pydantic model directly.
|
||||
# If we are using an older model or vLLM, we might need to manually parse the JSON.
|
||||
|
||||
# Using the newer 'beta.chat.completions.parse' for Pydantic support
|
||||
print(f"LLM Processor (Extract): Calling extraction for: {filtered_text}")
|
||||
try:
|
||||
completion = self.client.beta.chat.completions.parse(
|
||||
# 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=ExtractionResult,
|
||||
response_format={"type": "json_object"},
|
||||
extra_body={"include_reasoning": False},
|
||||
)
|
||||
return completion.choices[0].message.parsed
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user