File size: 6,384 Bytes
a52bae4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | # ============================================================================
# agent_py.py — Simple Python Agent backend (raw Mistral SDK tool-calling loop)
# ============================================================================
#
# CONTRACT: BACKEND_NAME, get_client, run, build_code_snippets
#
# PATTERN
# -------
# Classic tool-calling loop. The LLM sees the user's message plus a list
# of tool schemas. On each iteration it either:
# - emits tool calls (we run them and append results to the history), or
# - emits plain text (loop exits with that as the final reply).
#
# Bounded by MAX_AGENT_STEPS. No framework. Pure Python against the raw
# Mistral SDK.
# ============================================================================
import os
import json
# Defensive import — see agent_workflow.py for full explanation.
_Mistral = None
try:
from mistralai import Mistral as _Mistral # v1.x
except ImportError:
try:
from mistralai.client import Mistral as _Mistral # v2.x
except ImportError:
try:
from mistralai.client import MistralClient as _OldClient # v0.x
from mistralai.models.chat_completion import ChatMessage as _OldMsg
class _ChatShim:
def __init__(self, client):
self._client = client
def complete(self, model, messages, temperature=None,
max_tokens=None, tools=None):
msgs = [_OldMsg(role=m["role"], content=m.get("content", ""))
for m in messages]
return self._client.chat(
model=model, messages=msgs,
temperature=temperature, max_tokens=max_tokens,
tools=tools,
)
class _MistralV0Wrapper:
def __init__(self, api_key):
self._client = _OldClient(api_key=api_key)
self.chat = _ChatShim(self._client)
_Mistral = _MistralV0Wrapper
except ImportError as _e:
raise ImportError(
"mistralai package is missing or an unknown version. "
f"Last error: {_e}"
)
Mistral = _Mistral
from parameters import TEMPERATURE, MAX_TOKENS, MAX_AGENT_STEPS
from prompts import AGENT_SYSTEM
from tools import TOOL_FUNCTIONS, TOOL_SCHEMAS
import providers
BACKEND_NAME = "Simple Python Agent"
def get_client(api_key, provider="Mistral"):
return providers.get_llm_client(provider, api_key)
def _llm(client, messages, tools=None, provider="Mistral"):
model = providers.get_llm_model(provider)
return client.chat.complete(
model=model,
temperature=TEMPERATURE,
max_tokens=MAX_TOKENS,
messages=messages,
tools=tools,
).choices[0].message
def run(client, user_message, provider="Mistral"):
"""Tool-calling loop. LLM decides when to stop."""
messages = [
{"role": "system", "content": AGENT_SYSTEM},
{"role": "user", "content": user_message},
]
steps = []
tool_calls_made = []
for step_num in range(1, MAX_AGENT_STEPS + 1):
msg = _llm(client, messages, tools=TOOL_SCHEMAS, provider=provider)
messages.append(providers.serialize_assistant_message(msg, provider))
tool_calls = msg.tool_calls or []
if not tool_calls:
# No more tool calls — model has a final answer
steps.append({
"step": step_num, "type": "final", "tool": "-",
"args": "-", "result": msg.content or "",
})
return {
"reply": msg.content or "",
"steps": steps,
"extracted": {"tool_calls_made": tool_calls_made},
}
# Execute each tool call and append results
for tc in tool_calls:
name = tc.function.name
args_raw = tc.function.arguments
args = json.loads(args_raw) if isinstance(args_raw, str) else args_raw
result = TOOL_FUNCTIONS[name](**args)
steps.append({
"step": step_num, "type": "tool_call", "tool": name,
"args": json.dumps(args), "result": str(result),
})
tool_calls_made.append({"tool": name, "args": args, "result": result})
messages.append(providers.serialize_tool_result(tc, name, result, provider))
steps.append({
"step": MAX_AGENT_STEPS, "type": "limit", "tool": "-",
"args": "-", "result": "max steps reached",
})
return {
"reply": "(max agent steps reached)",
"steps": steps,
"extracted": {"tool_calls_made": tool_calls_made},
}
def build_code_snippets(user_message, steps):
lines = [
"# Backend: Simple Python Agent",
"# Raw Mistral SDK tool-calling loop. No framework.",
f"# User message: {user_message}",
"",
"messages = [",
" {'role': 'system', 'content': AGENT_SYSTEM},",
f" {{'role': 'user', 'content': {user_message!r}}},",
"]",
"",
"for step in range(1, MAX_AGENT_STEPS + 1):",
" msg = client.chat.complete(",
" model=MODEL, messages=messages, tools=TOOL_SCHEMAS",
" ).choices[0].message",
" messages.append(msg.model_dump(exclude_none=True))",
"",
" if not msg.tool_calls:",
" break # plain-text reply means we are done",
"",
" for tc in msg.tool_calls:",
" name = tc.function.name",
" args = json.loads(tc.function.arguments)",
" result = TOOL_FUNCTIONS[name](**args)",
" messages.append({",
" 'role': 'tool', 'name': name,",
" 'content': result, 'tool_call_id': tc.id,",
" })",
"",
"# ---------- actual step log ----------",
]
for s in steps:
lines.append(f"# Step {s['step']} [{s['type']}] tool={s['tool']}")
lines.append(f"# args: {s['args']}")
lines.append(f"# result: {s['result']}")
return "\n".join(lines)
|