RUN YOUR AGENT AS A SUBAGENT IN HERMES

name: stock-agent
version: "1.0"
description: Gives the main model the stock lookup agent on Oumi as a tool

import json
import urllib.error
import urllib.parse
import urllib.request

API = "https://api.example-stocks.com"

# Copies of the Oumi tool definitions. A name, parameter, or description that
# differs from training makes the model call tools that do not exist.
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "search_company",
            "description": "Find a company's stock ticker symbol by company name. Call this before get_stock_price.",
            "parameters": {
                "type": "object",
                "properties": {"name": {"type": "string", "description": "Company name, e.g. 'Apple'."}},
                "required": ["name"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_stock_price",
            "description": "Get the latest stock price for a ticker symbol returned by search_company.",
            "parameters": {
                "type": "object",
                "properties": {"ticker": {"type": "string", "description": "Stock ticker symbol, e.g. 'AAPL'."}},
                "required": ["ticker"],
            },
        },
    },
]


def run_tool(name: str, args: dict) -> str:
    if name == "search_company":
        return search_company(args["name"])
    if name == "get_stock_price":
        return get_stock_price(args["ticker"])
    raise ValueError(f"Unknown tool: {name}")


def search_company(name: str) -> str:
    found = get_json(f"{API}/search?q={urllib.parse.quote(name)}")
    if found is None:
        return f'No company found matching "{name}".'
    return f"Found {found['name']} (ticker: {found['ticker']})."


def get_stock_price(ticker: str) -> str:
    quote = get_json(f"{API}/quote/{urllib.parse.quote(ticker)}")
    if quote is None:
        return f'No price found for "{ticker}".'
    return f"{ticker}: {quote['price']} {quote['currency']}."


# A 404 means "no match" and returns None. Any other failure raises.
def get_json(url: str) -> dict | None:
    try:
        with urllib.request.urlopen(url, timeout=30) as res:
            return json.load(res)
    except urllib.error.HTTPError as error:
        if error.code == 404:
            return None
        raise

import json
import urllib.error
import urllib.request
from pathlib import Path

from .tools import TOOLS, run_tool

ENDPOINT = "https://inference.oumi.ai/v1/chat/completions"
MODEL = "projects/{project_id}/deployments/{deployment_id}"

# Stops a model that keeps calling tools. Set it above the number of tool
# calls that your longest training conversation makes.
MAX_TOOL_STEPS = 8

SYSTEM_PROMPT = (Path(__file__).parent / "system-prompt.md").read_text()


class StockAgent:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.history = [{"role": "system", "content": SYSTEM_PROMPT}]

    # A turn that fails part way leaves tool calls without results, and your
    # deployment rejects a conversation like that. Removing the failed turn
    # keeps the next question working.
    def ask(self, question: str) -> str:
        start = len(self.history)
        try:
            return self._turn(question)
        except Exception:
            del self.history[start:]
            raise

    def _turn(self, question: str) -> str:
        self.history.append({"role": "user", "content": question})

        for _ in range(MAX_TOOL_STEPS):
            reply = self._complete()
            calls = reply.get("tool_calls") or []

            if not calls:
                self.history.append({"role": "assistant", "content": reply["content"]})
                return reply["content"]

            self.history.append({"role": "assistant", "content": reply.get("content") or "", "tool_calls": calls})
            for call in calls:
                self.history.append({"role": "tool", "tool_call_id": call["id"], "content": run_tool_call(call)})

        raise RuntimeError(f"The stock agent made {MAX_TOOL_STEPS} rounds of tool calls without an answer.")

    def _complete(self) -> dict:
        request = urllib.request.Request(
            ENDPOINT,
            data=json.dumps({"model": MODEL, "messages": self.history, "tools": TOOLS}).encode(),
            headers={"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"},
        )
        try:
            with urllib.request.urlopen(request, timeout=120) as res:
                return json.load(res)["choices"][0]["message"]
        except urllib.error.HTTPError as error:
            raise RuntimeError(f"Oumi returned HTTP {error.code}: {error.read().decode()}") from error


# A failed tool call goes back to the model as its result, so the model can
# tell the user instead of inventing a value.
def run_tool_call(call: dict) -> str:
    try:
        return run_tool(call["function"]["name"], json.loads(call["function"]["arguments"]))
    except Exception as error:
        return f"Error: {error}"

from pathlib import Path

from .agent import StockAgent

API_KEY_FILE = Path.home() / ".config/oumi/api.key"

SCHEMA = {
    "name": "ask_stock_agent",
    "description": (
        "Ask the stock agent for the latest stock price of one or more companies. It finds the "
        "ticker symbols itself and answers in prose. It remembers earlier questions in this "
        "session, so a follow-up can refer to a company it already found. Set reset to true to "
        "start a new conversation with it."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "question": {"type": "string", "description": "The question, in plain English."},
            "reset": {"type": "boolean", "description": "Forget earlier questions before this one."},
        },
        "required": ["question"],
    },
}

# One subagent per Hermes session. It keeps its own history, so a follow-up
# question can refer to a company that it already found.
agents: dict[str | None, StockAgent] = {}


def ask_stock_agent(args: dict, session_id: str | None = None, **kwargs) -> str:
    if session_id not in agents or args.get("reset"):
        agents[session_id] = StockAgent(API_KEY_FILE.read_text().strip())
    return agents[session_id].ask(args["question"])


def register(ctx):
    ctx.register_tool(name="ask_stock_agent", toolset="stock_agent", schema=SCHEMA, handler=ask_stock_agent)