add agent.py

agent.py

@@ -0,0 +1,64 @@
+"""
+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"]