42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
"""Small LangChain agent used by the stream-mode console demo."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain_core.tools import tool
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
|
|
@tool
|
|
def shopping_list(query: str) -> str:
|
|
"""Return a compact shopping list for a grocery-related request."""
|
|
|
|
normalized = query.lower()
|
|
if "молоко" in normalized or "хлеб" in normalized or "яблок" in normalized:
|
|
return "молоко, хлеб, яблоки"
|
|
return "молоко, хлеб, яблоки, чай"
|
|
|
|
|
|
def build_model() -> ChatOpenAI:
|
|
"""Create an OpenAI-compatible chat model from environment variables."""
|
|
|
|
return ChatOpenAI(
|
|
model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
|
|
base_url=os.getenv("OPENAI_BASE_URL") or None,
|
|
api_key=os.getenv("OPENAI_API_KEY", "not-needed"),
|
|
temperature=0,
|
|
streaming=True,
|
|
)
|
|
|
|
|
|
agent = create_agent(
|
|
model=build_model(),
|
|
tools=[shopping_list],
|
|
system_prompt=(
|
|
"Ты полезный консольный AI-агент. Отвечай кратко. "
|
|
"Если пользователь просит список покупок, используй shopping_list."
|
|
),
|
|
)
|