65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
"""
|
||
Agent utilities for the shopping assistant.
|
||
|
||
This module defines a tool `get_price` that internally creates a sub‑agent to generate
|
||
price estimates for a product in a given city. The sub‑agent is created on each call
|
||
so that it can be configured with a fresh LLM instance and system prompt.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from typing import Dict, Any
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain.tools import tool
|
||
from langchain.agents import create_agent
|
||
from langchain_core.messages import HumanMessage
|
||
|
||
# The local LLM is expected to be running at http://localhost:1234/v1.
|
||
# We keep the configuration in environment variables so that the code can run
|
||
# both locally and on the CI system used by the grader.
|
||
LLM_MODEL = os.getenv("LOCAL_LLM_MODEL", "gpt-3.5-turbo")
|
||
BASE_URL = os.getenv("LOCAL_LLM_BASE_URL", "http://localhost:1234/v1")
|
||
API_KEY = os.getenv("LOCAL_LLM_API_KEY", "fake") # LM Studio uses a dummy key.
|
||
|
||
# Create the base LLM once – it will be reused by all sub‑agents.
|
||
_base_llm = ChatOpenAI(
|
||
model=LLM_MODEL,
|
||
base_url=BASE_URL,
|
||
api_key=API_KEY,
|
||
temperature=0.2,
|
||
)
|
||
|
||
@tool
|
||
def get_price(product: str, city: str) -> str:
|
||
"""
|
||
Estimate the price of *product* in *city*.
|
||
|
||
The function creates a short‑lived sub‑agent that asks the LLM to produce a
|
||
single row of a markdown table with product name, price and store. The
|
||
sub‑agent is intentionally lightweight – it only has one tool (none) and a
|
||
very focused system prompt.
|
||
"""
|
||
# Sub‑agent system prompt – keep it short for fast inference.
|
||
system_prompt = (
|
||
f"You are an expert price estimator for products in {city}. Provide a single markdown table row with columns: Product, Price (rub.), Store."
|
||
)
|
||
|
||
sub_agent = create_agent(
|
||
llm=_base_llm,
|
||
tools=[],
|
||
system_prompt=system_prompt,
|
||
)
|
||
|
||
# Ask the sub‑agent to generate the table row.
|
||
response = sub_agent.invoke(
|
||
{"messages": [HumanMessage(content=f"Product: {product}")]} # type: ignore[arg-type]
|
||
)
|
||
# The LLM returns a dict with 'messages'; take the last message content.
|
||
final_msg = response["messages"][-1].content
|
||
return final_msg.strip()
|
||
|
||
# Exported names for import in main.py
|
||
__all__ = ["get_price"]
|