|
|
| 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 |
| } |
|
|
