File size: 1,520 Bytes
1870dee df07284 1870dee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import json
from typing import Any
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage
from dotenv import load_dotenv
from .state import AgentState, parameter_catalog
load_dotenv()
class ExtractedParameter(BaseModel):
name: str = Field(
description="Parameter name from the global catalog"
)
value: str | int | float | bool | None = Field(
description="Extracted parameter value"
)
class ExtractionResult(BaseModel):
parameters: list[ExtractedParameter] = Field(
default_factory=list
)
llm = ChatOpenAI(model="gpt-4o")
extractor_llm = llm.with_structured_output(
ExtractionResult
)
def extract_parameters(state: AgentState):
system_prompt = f"""
You are a parameter extraction agent.
Extract parameters from the user conversation.
Available parameter catalog:
{json.dumps(parameter_catalog)}
Rules:
- Only extract parameters defined in the catalog.
- Use conversation context to understand references.
- Do not invent missing values.
- Return only new or corrected values.
- Do not decide which tool to execute.
"""
result = extractor_llm.invoke([
SystemMessage(content=system_prompt),
*state["messages"]
])
extracted = {
param.name: param.value
for param in result.parameters
if param.name in parameter_catalog
}
return {
"global_params": extracted
}
|