Files
TraderAI/traderai/agent.py
T
2026-05-07 21:20:43 -04:00

630 lines
28 KiB
Python

from __future__ import annotations
import json
import re
from collections.abc import AsyncIterator
from typing import Any
import httpx
from tzlocal import get_localzone
from traderai.memory import DEFAULT_THREAD_ID, MemoryStore, iso_now, iso_now_in_zone, time_since
from traderai.tools import ToolRegistry
SYSTEM_PROMPT = """You are TraderAI, a local assistant for UEX marketplace work.
Use tools when the user asks about UEX data, open/current listings, active negotiations, unread notifications, messages, offers, or posting ads.
UEX credentials are configured server-side when available. Never ask the user to provide UEX_SECRET_KEY or UEX_BEARER_TOKEN in chat; call the authenticated UEX tool and only mention credential configuration if the tool returns an authentication error.
Use the specific UEX tool for the needed endpoint, such as get_uex_commodities_prices or get_uex_vehicles. Use fields, limit, and summary mode so tool results stay compact.
When the user asks for history, trends, changes over time, or past prices, prefer the summarize_uex_*_history tools when available; use search_uex_api_index(history_only=true) if you need to discover history endpoints.
Use SCMDB tools when the user asks about Star Citizen missions/contracts, mission rewards, payouts, reputation gains, item rewards, blueprint rewards, or hauling mission cargo. Prefer SCMDB live data unless the user asks for PTU or a specific game version.
Prefer open and current UEX marketplace information. Do not use historical sale data, completed sale records, or sale/average-history information unless the user explicitly asks for historical sales.
Treat UEX marketplace prices as in-game aUEC/UEC credits, never real-world dollars, unless the user explicitly says otherwise.
For marketplace writes, draft the exact pending action and tell the user what will be sent; never claim it was sent until approval succeeds.
When a scheduled wake job fires, always write a concise Inbox-ready result that says what you checked, the key findings, and the suggested next action.
Keep prices, listing ids, slugs, users, and UEX status codes precise. If data is missing, say what you need next."""
class OllamaAgent:
def __init__(
self,
base_url: str,
model: str,
tools: ToolRegistry,
memory: MemoryStore | None = None,
user_name: str | None = None,
num_ctx: int | None = None,
) -> None:
self.base_url = base_url.rstrip("/")
self.model = model
self.tools = tools
self.memory = memory
self.user_name = user_name
self.num_ctx = num_ctx
self.thread_messages: dict[str, list[dict[str, Any]]] = {}
async def health(self) -> dict[str, Any]:
try:
async with httpx.AsyncClient(timeout=3) as client:
response = await client.get(f"{self.base_url}/api/tags")
response.raise_for_status()
body = response.json()
except (httpx.HTTPError, ValueError) as exc:
return {
"online": False,
"model": self.model,
"base_url": self.base_url,
"message": f"Ollama is offline or unreachable at {self.base_url}. Open the Ollama tab and use the recommended action.",
"detail": str(exc),
}
models = [model.get("name") or model.get("model") for model in body.get("models", [])]
return {
"online": True,
"model": self.model,
"base_url": self.base_url,
"model_available": self.model in models,
"models": models,
"message": "Ollama is online.",
}
async def ensure_available(self) -> None:
health = await self.health()
if not health["online"]:
raise OllamaUnavailable(health["message"])
async def chat(self, content: str, thread_id: str | None = DEFAULT_THREAD_ID) -> dict[str, Any]:
await self.ensure_available()
resolved_thread_id = self._thread_id(thread_id)
messages = self._messages_for_thread(resolved_thread_id)
previous_interaction = self.memory.last_interaction(resolved_thread_id) if self.memory else None
if self.memory:
self.memory.add_conversation("user", content, resolved_thread_id)
await self._title_first_message(resolved_thread_id, content, previous_interaction)
messages.append({"role": "user", "content": content})
last_tool_results: list[dict[str, Any]] = []
for _ in range(5):
try:
response = await self._ollama_chat(
content,
messages,
previous_interaction=previous_interaction,
thread_id=resolved_thread_id,
)
except Exception as exc:
if not last_tool_results:
raise
answer = self._tool_result_fallback(
last_tool_results,
f"The local model stopped after the tool call: {exc}",
)
messages.append({"role": "assistant", "content": answer})
if self.memory:
self.memory.add_conversation("assistant", answer, resolved_thread_id)
return {"message": answer, "pending_actions": self._pending_payloads(), "thread_id": resolved_thread_id}
message = response.get("message") or {}
tool_calls = message.get("tool_calls") or []
if not tool_calls:
answer = message.get("content", "")
if not answer.strip():
answer = self._empty_response_fallback(last_tool_results)
messages.append({"role": "assistant", "content": answer})
if self.memory:
self.memory.add_conversation("assistant", answer, resolved_thread_id)
return {"message": answer, "pending_actions": self._pending_payloads(), "thread_id": resolved_thread_id}
messages.append(message)
for call in tool_calls:
name, arguments = self._extract_call(call)
result = await self.tools.execute(name, arguments)
last_tool_results.append({"tool": name, "result": result})
messages.append({"role": "tool", "tool_name": name, "content": json.dumps(result)})
fallback = "I hit the tool-call limit while working on that. Try narrowing the request or approve any pending action first."
messages.append({"role": "assistant", "content": fallback})
if self.memory:
self.memory.add_conversation("assistant", fallback, resolved_thread_id)
return {"message": fallback, "pending_actions": self._pending_payloads(), "thread_id": resolved_thread_id}
async def chat_events(self, content: str, thread_id: str | None = DEFAULT_THREAD_ID) -> AsyncIterator[dict[str, Any]]:
health = await self.health()
if not health["online"]:
yield {"type": "warning", "message": health["message"]}
yield {"type": "done", "pending_actions": self._pending_payloads()}
return
resolved_thread_id = self._thread_id(thread_id)
messages = self._messages_for_thread(resolved_thread_id)
previous_interaction = self.memory.last_interaction(resolved_thread_id) if self.memory else None
if self.memory:
self.memory.add_conversation("user", content, resolved_thread_id)
await self._title_first_message(resolved_thread_id, content, previous_interaction)
messages.append({"role": "user", "content": content})
yield {"type": "status", "message": "Thinking"}
last_tool_results: list[dict[str, Any]] = []
for _ in range(5):
assistant_message: dict[str, Any] = {"role": "assistant", "content": ""}
tool_calls: list[dict[str, Any]] = []
try:
async for event in self._ollama_chat_stream(
content,
messages,
previous_interaction=previous_interaction,
thread_id=resolved_thread_id,
):
message = event.get("message") or {}
chunk = message.get("content") or ""
if chunk:
assistant_message["content"] += chunk
yield {"type": "token", "content": chunk}
if message.get("tool_calls"):
tool_calls.extend(message["tool_calls"])
if event.get("done"):
metrics = self._stream_metrics(event)
if metrics:
yield {"type": "metrics", **metrics}
except Exception as exc:
if not last_tool_results:
yield {"type": "warning", "message": f"Chat failed before any tool result was available: {exc}"}
yield {"type": "done", "pending_actions": self._pending_payloads(), "thread_id": resolved_thread_id}
return
fallback = self._tool_result_fallback(
last_tool_results,
f"The local model stopped after the tool call: {exc}",
)
assistant_message["content"] = fallback
messages.append(assistant_message)
if self.memory:
self.memory.add_conversation("assistant", fallback, resolved_thread_id)
yield {"type": "token", "content": fallback}
yield {"type": "done", "pending_actions": self._pending_payloads(), "thread_id": resolved_thread_id}
return
if not tool_calls:
if not assistant_message.get("content", "").strip():
fallback = self._empty_response_fallback(last_tool_results)
assistant_message["content"] = fallback
yield {"type": "token", "content": fallback}
messages.append(assistant_message)
if self.memory:
self.memory.add_conversation("assistant", assistant_message.get("content", ""), resolved_thread_id)
yield {"type": "done", "pending_actions": self._pending_payloads(), "thread_id": resolved_thread_id}
return
assistant_message["tool_calls"] = tool_calls
messages.append(assistant_message)
for call in tool_calls:
name, arguments = self._extract_call(call)
yield {"type": "status", "message": self._tool_status(name)}
result = await self.tools.execute(name, arguments)
last_tool_results.append({"tool": name, "result": result})
messages.append({"role": "tool", "tool_name": name, "content": json.dumps(result)})
yield {"type": "status", "message": "Writing response"}
fallback = "I hit the tool-call limit while working on that. Try narrowing the request or approve any pending action first."
messages.append({"role": "assistant", "content": fallback})
if self.memory:
self.memory.add_conversation("assistant", fallback, resolved_thread_id)
yield {"type": "token", "content": fallback}
yield {"type": "done", "pending_actions": self._pending_payloads(), "thread_id": resolved_thread_id}
async def generate_wake_response(self, wake_message: str) -> str:
await self.ensure_available()
messages = self._messages_for_thread("wake")
previous_interaction = self.memory.last_interaction("wake") if self.memory else None
messages.append({"role": "user", "content": wake_message})
last_tool_results: list[dict[str, Any]] = []
for _ in range(5):
try:
response = await self._ollama_chat(
wake_message,
messages,
previous_interaction=previous_interaction,
thread_id="wake",
)
except Exception as exc:
if not last_tool_results:
raise
content = self._tool_result_fallback(
last_tool_results,
f"The local model stopped after the wake-job tool call: {exc}",
)
messages.append({"role": "assistant", "content": content})
if self.memory:
self.memory.add_conversation("system", wake_message, "wake")
self.memory.add_conversation("assistant", content, "wake")
return content
message = response.get("message") or {}
tool_calls = message.get("tool_calls") or []
if not tool_calls:
content = message.get("content", "")
if not content.strip():
content = self._empty_response_fallback(last_tool_results)
messages.append({"role": "assistant", "content": content})
if self.memory:
self.memory.add_conversation("system", wake_message, "wake")
self.memory.add_conversation("assistant", content, "wake")
return content
messages.append(message)
for call in tool_calls:
name, arguments = self._extract_call(call)
result = await self.tools.execute(name, arguments)
last_tool_results.append({"tool": name, "result": result})
messages.append({"role": "tool", "tool_name": name, "content": json.dumps(result)})
content = "I hit the tool-call limit while running this scheduled wake job. Check the job prompt or pending approvals."
messages.append({"role": "assistant", "content": content})
if self.memory:
self.memory.add_conversation("system", wake_message, "wake")
self.memory.add_conversation("assistant", content, "wake")
return content
async def _ollama_chat(
self,
query: str = "",
messages: list[dict[str, Any]] | None = None,
previous_interaction: dict[str, Any] | None = None,
thread_id: str | None = DEFAULT_THREAD_ID,
) -> dict[str, Any]:
async with httpx.AsyncClient(timeout=120) as client:
response = await client.post(
f"{self.base_url}/api/chat",
json={
"model": self.model,
"messages": self._messages_with_context(
query,
messages or self._messages_for_thread(thread_id),
previous_interaction=previous_interaction,
thread_id=thread_id,
),
"tools": self.tools.schemas,
"options": self._ollama_options(),
"stream": False,
},
)
response.raise_for_status()
return response.json()
async def _ollama_chat_stream(
self,
query: str = "",
messages: list[dict[str, Any]] | None = None,
previous_interaction: dict[str, Any] | None = None,
thread_id: str | None = DEFAULT_THREAD_ID,
) -> AsyncIterator[dict[str, Any]]:
async with httpx.AsyncClient(timeout=120) as client:
async with client.stream(
"POST",
f"{self.base_url}/api/chat",
json={
"model": self.model,
"messages": self._messages_with_context(
query,
messages or self._messages_for_thread(thread_id),
previous_interaction=previous_interaction,
thread_id=thread_id,
),
"tools": self.tools.schemas,
"options": self._ollama_options(),
"stream": True,
},
) as response:
response.raise_for_status()
async for line in response.aiter_lines():
if line:
yield json.loads(line)
def _messages_with_context(
self,
query: str,
messages: list[dict[str, Any]],
previous_interaction: dict[str, Any] | None = None,
thread_id: str | None = DEFAULT_THREAD_ID,
) -> list[dict[str, Any]]:
context = self._runtime_context(query, previous_interaction=previous_interaction, thread_id=thread_id)
if not context:
return messages
return [messages[0], {"role": "system", "content": context}, *messages[1:]]
def _runtime_context(
self,
query: str,
previous_interaction: dict[str, Any] | None = None,
thread_id: str | None = DEFAULT_THREAD_ID,
) -> str:
local_zone = get_localzone()
parts = [
f"Current local date/time: {iso_now()} UTC; {iso_now_in_zone(local_zone)} {local_zone}.",
]
uex = getattr(self.tools, "uex", None)
if uex:
auth_methods = []
if uex.secret_key:
auth_methods.append("secret key")
if uex.bearer_token:
auth_methods.append("bearer token")
if auth_methods:
parts.append(
"UEX API authentication is configured server-side with "
+ " and ".join(auth_methods)
+ "; use authenticated UEX tools directly and do not ask for tokens."
)
else:
parts.append("UEX API authentication is not configured server-side.")
if self.user_name:
parts.append(f"Known user name/handle: {self.user_name}.")
if self.memory is None:
return "\n".join(parts)
profile = self.memory.get_profile()
if profile:
identity = self._profile_identity(profile)
if identity:
parts.append(identity)
parts.append(f"Known user profile JSON: {json.dumps(self._profile_for_prompt(profile), ensure_ascii=True)}.")
last = previous_interaction if previous_interaction is not None else self.memory.last_interaction(thread_id)
if last:
parts.append(
f"Previous interaction before this message: {last['created_at']} "
f"({time_since(last['created_at'])}, role {last['role']})."
)
else:
parts.append("Previous interaction before this message: none recorded.")
memories = self.memory.recall(query, limit=6)
if memories:
memory_text = "\n".join(
f"- [{item['kind']}, importance {item['importance']}] {item['content']}"
for item in memories
)
parts.append(f"Relevant long-term memories:\n{memory_text}")
recent = self.memory.recent_conversation(limit=6, thread_id=thread_id)
if recent:
recent_text = "\n".join(
f"- {item['created_at']} {item['role']}: {item['content'][:500]}"
for item in recent
)
parts.append(f"Recent conversation excerpts from this chat:\n{recent_text}")
return "\n".join(parts)
def _messages_for_thread(self, thread_id: str | None) -> list[dict[str, Any]]:
resolved_thread_id = self._thread_id(thread_id)
if resolved_thread_id not in self.thread_messages:
messages: list[dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
if self.memory:
self.memory.ensure_thread(resolved_thread_id)
for item in self.memory.recent_conversation(limit=30, thread_id=resolved_thread_id):
role = item.get("role")
if role in {"user", "assistant"} and item.get("content"):
messages.append({"role": role, "content": item["content"]})
self.thread_messages[resolved_thread_id] = messages
return self.thread_messages[resolved_thread_id]
async def _title_first_message(
self,
thread_id: str,
first_message: str,
previous_interaction: dict[str, Any] | None,
) -> None:
if self.memory is None or previous_interaction is not None:
return
thread = self.memory.get_thread(thread_id)
if not thread or thread.get("title") != "New chat":
return
title = await self._generate_chat_title(first_message)
self.memory.rename_thread(thread_id, title or MemoryStore._thread_title(first_message))
async def _generate_chat_title(self, first_message: str) -> str:
prompt = (
"Create a concise chat title for this first user message. "
"Use 2 to 6 words. No quotes, no punctuation at the end, no preamble.\n\n"
f"Message: {first_message[:800]}"
)
try:
async with httpx.AsyncClient(timeout=20) as client:
response = await client.post(
f"{self.base_url}/api/chat",
json={
"model": self.model,
"messages": [
{"role": "system", "content": "You write short chat titles."},
{"role": "user", "content": prompt},
],
"options": self._ollama_options(),
"stream": False,
},
)
response.raise_for_status()
message = response.json().get("message") or {}
return self._clean_generated_title(message.get("content", ""))
except Exception:
return ""
@staticmethod
def _thread_id(thread_id: str | None) -> str:
return (thread_id or DEFAULT_THREAD_ID).strip() or DEFAULT_THREAD_ID
@staticmethod
def _clean_generated_title(title: str) -> str:
text = re.sub(r"[\r\n]+", " ", title).strip().strip('"').strip("'")
text = re.sub(r"^(title|chat title)\s*:\s*", "", text, flags=re.IGNORECASE).strip()
text = text.rstrip(".!?;:-").strip()
if not text:
return ""
words = text.split()
if len(words) > 8:
text = " ".join(words[:8])
return text[:64]
def _pending_payloads(self) -> list[dict[str, Any]]:
return [
{
"id": action.id,
"label": action.label,
"method": action.method,
"endpoint": action.endpoint,
"payload": action.payload,
}
for action in self.tools.pending_actions.values()
]
def _ollama_options(self) -> dict[str, Any]:
if not self.num_ctx:
return {}
return {"num_ctx": self.num_ctx}
@staticmethod
def _empty_response_fallback(tool_results: list[dict[str, Any]]) -> str:
if not tool_results:
return "I did not get a usable response from the local model. Please try again, or narrow the request a bit."
return OllamaAgent._tool_result_fallback(
tool_results,
"I completed the tool call, but the local model did not write a final answer.",
)
@staticmethod
def _tool_result_fallback(tool_results: list[dict[str, Any]], reason: str) -> str:
last = tool_results[-1]
text = json.dumps(last, indent=2, ensure_ascii=True)
if len(text) > 1800:
text = text[:1800] + "\n..."
return (
f"{reason} "
"Here is the last tool result so you are not left staring at a blank response:\n\n"
f"```json\n{text}\n```"
)
@staticmethod
def _tool_status(name: str) -> str:
if name.startswith("get_uex_"):
return f"Fetching UEX {name.removeprefix('get_uex_')}"
if name.startswith("draft_uex_"):
return f"Drafting UEX {name.removeprefix('draft_uex_')} for approval"
if name.startswith("delete_uex_"):
return f"Drafting UEX {name.removeprefix('delete_uex_')} delete for approval"
labels = {
"search_uex_api_index": "Searching UEX API index",
"summarize_uex_commodity_price_history": "Summarizing commodity price history",
"summarize_uex_marketplace_price_history": "Summarizing marketplace price history",
"summarize_uex_currency_index_history": "Summarizing currency index history",
"list_scmdb_versions": "Checking SCMDB versions",
"search_scmdb_missions": "Searching SCMDB missions",
"get_scmdb_mission_rewards": "Fetching SCMDB mission rewards",
"uex_api_catalog": "Checking UEX API catalog",
"uex_get": "Fetching UEX data",
"uex_draft_post": "Drafting UEX write for approval",
"uex_draft_delete": "Drafting UEX delete for approval",
"search_marketplace_listings": "Searching UEX listings",
"get_marketplace_listing": "Fetching listing details",
"list_marketplace_negotiations": "Checking negotiations",
"get_negotiation_messages": "Reading negotiation messages",
"draft_negotiation_message": "Drafting message for approval",
"draft_marketplace_listing": "Drafting listing for approval",
"check_uex_notifications": "Checking UEX notifications",
}
return labels.get(name, f"Running {name}")
@staticmethod
def _stream_metrics(event: dict[str, Any]) -> dict[str, Any]:
prompt_tokens = int(event.get("prompt_eval_count") or 0)
prompt_duration = int(event.get("prompt_eval_duration") or 0)
output_tokens = int(event.get("eval_count") or 0)
output_duration = int(event.get("eval_duration") or 0)
def rate(tokens: int, duration_ns: int) -> float | None:
if not tokens or not duration_ns:
return None
return tokens / (duration_ns / 1_000_000_000)
return {
"reading_tokens": prompt_tokens,
"reading_tokens_per_second": rate(prompt_tokens, prompt_duration),
"writing_tokens": output_tokens,
"writing_tokens_per_second": rate(output_tokens, output_duration),
}
@staticmethod
def _profile_identity(profile: dict[str, Any]) -> str:
user = profile.get("uex_user")
if not isinstance(user, dict):
configured = profile.get("configured_name")
return f"You are speaking with {configured}." if configured else ""
username = user.get("username") or user.get("user_username")
name = user.get("name")
fields = []
if username and name and username != name:
fields.append(f"You are speaking with UEX user {username} ({name}).")
elif username or name:
fields.append(f"You are speaking with UEX user {username or name}.")
details = []
for key, label in [
("timezone", "timezone"),
("language", "preferred language"),
("specializations", "specializations"),
("languages", "languages"),
("archetypes", "archetypes"),
]:
value = user.get(key)
if value:
details.append(f"{label}: {value}")
if details:
fields.append("UEX profile details: " + "; ".join(details) + ".")
return " ".join(fields)
@staticmethod
def _profile_for_prompt(profile: dict[str, Any]) -> dict[str, Any]:
user = profile.get("uex_user")
if not isinstance(user, dict):
return profile
useful_user_fields = [
"id",
"name",
"username",
"avatar",
"bio",
"website_url",
"timezone",
"language",
"day_availability",
"time_availability",
"specializations",
"languages",
"archetypes",
"is_datarunner",
"is_staff",
"is_away_game",
"date_rsi_verified",
"date_twitch_verified",
]
prompt_profile = dict(profile)
prompt_profile["uex_user"] = {
key: user[key]
for key in useful_user_fields
if key in user and user[key] not in (None, "")
}
return prompt_profile
@staticmethod
def _extract_call(call: dict[str, Any]) -> tuple[str, dict[str, Any]]:
function = call.get("function") or {}
name = function.get("name") or call.get("name")
arguments = function.get("arguments") or call.get("arguments") or {}
if isinstance(arguments, str):
arguments = json.loads(arguments or "{}")
return name, arguments
class OllamaUnavailable(RuntimeError):
pass