Something went wrong. Try again.
Stateful AI assistants running off Claude/Gemini CLI subscription
Something went wrong. Try again.
9.3 kB · 230 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231import asyncioimport jsonimport osimport sysimport timefrom datetime import datetimefrom pathlib import Path
from google import genaifrom google.genai import typesfrom mcp import ClientSession, StdioServerParametersfrom mcp.client.stdio import stdio_client
# Import helpers from invoke_dot# We assume invoke_dot.py is in the same directory (poc/)sys.path.insert(0, str(Path(__file__).parent))from invoke_dot import HOME, MESSAGES_LOG, log_event, write_session_log, read_new_messages
INTERACTIONS_CACHE = HOME / "logs" / ".gemini_interactions.json"
def mcp_to_gemini_schema(mcp_schema: dict) -> types.Schema: """Convert MCP JSON Schema to Gemini Schema type (for generate_content API).""" if not mcp_schema: return types.Schema(type=types.Type.OBJECT, properties={})
schema_copy = mcp_schema.copy() for field in ["$schema", "additionalProperties", "$defs", "$ref", "title"]: schema_copy.pop(field, None)
type_map = { "object": types.Type.OBJECT, "string": types.Type.STRING, "integer": types.Type.INTEGER, "number": types.Type.NUMBER, "boolean": types.Type.BOOLEAN, "array": types.Type.ARRAY, } gemini_type = type_map.get(schema_copy.get("type"), types.Type.OBJECT)
properties = None if "properties" in schema_copy: properties = {k: mcp_to_gemini_schema(v) for k, v in schema_copy["properties"].items()}
items = None if "items" in schema_copy: items = mcp_to_gemini_schema(schema_copy["items"])
return types.Schema( type=gemini_type, properties=properties, required=schema_copy.get("required"), description=schema_copy.get("description"), enum=schema_copy.get("enum"), items=items, )
_UNSUPPORTED_SCHEMA_FIELDS = { "$schema", "additionalProperties", "$defs", "$ref", "title", "default", "format", "examples", "$comment", "if", "then", "else", "not", "anyOf", "oneOf", "allOf",}
def _clean_schema_dict(schema: dict) -> dict: """Strip unsupported JSON Schema fields, returning a plain dict.
Used for the Interactions API which requires plain JSON, not SDK types. Gemini only supports: type, properties, required, description, enum, items. """ if not schema: return {"type": "object", "properties": {}} result = {k: v for k, v in schema.items() if k not in _UNSUPPORTED_SCHEMA_FIELDS} if "properties" in result: result["properties"] = {k: _clean_schema_dict(v) for k, v in result["properties"].items()} if "items" in result: result["items"] = _clean_schema_dict(result["items"]) # Omit empty required arrays — some API validators reject them if "required" in result and result["required"] == []: del result["required"] return result
async def _run_gemini_loop(prompt: str, model: str, thinking_level: str, conversation_id: str | None): """Inner loop that handles MCP session and Gemini Interactions API calls.""" server_params = StdioServerParameters( command="uv", args=["run", "python", "dot_mcp_server.py"], env={**os.environ, "DOT_HOME": str(HOME)} )
async with stdio_client(server_params) as (read, write): async with ClientSession(read, write) as session: await session.initialize() # List tools and convert to Gemini format tools_result = await session.list_tools() tools = [ { "type": "function", "name": tool.name, "description": tool.description, "parameters": _clean_schema_dict(tool.inputSchema), } for tool in tools_result.tools ]
client = genai.Client( api_key=os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY") )
# Handle stateful conversation — only chain Discord conversations, # not scheduler ticks (each tick is independent; chaining causes # unbounded server-side history growth and timeouts). is_stateful = conversation_id and not conversation_id.startswith("scheduler:") prev_id = None interactions_map = {} if is_stateful: INTERACTIONS_CACHE.parent.mkdir(parents=True, exist_ok=True) if INTERACTIONS_CACHE.exists(): try: interactions_map = json.loads(INTERACTIONS_CACHE.read_text(encoding="utf-8")) prev_id = interactions_map.get(conversation_id) except Exception: pass
current_input = prompt
while True: create_kwargs: dict = dict(model=model, input=current_input, tools=tools) if prev_id: create_kwargs["previous_interaction_id"] = prev_id if thinking_level: create_kwargs["generation_config"] = {"thinking_level": thinking_level} print(f"[orchestrator] Gemini call: prev_id={bool(prev_id)}, " f"input_type={'str' if isinstance(current_input, str) else f'list[{len(current_input)}]'}", file=sys.stderr) interaction = await client.aio.interactions.create(**create_kwargs) prev_id = interaction.id
# Update interaction cache (stateful conversations only) if is_stateful: interactions_map[conversation_id] = interaction.id INTERACTIONS_CACHE.write_text(json.dumps(interactions_map), encoding="utf-8")
if interaction.status == "completed": break elif interaction.status == "requires_action": # Execute MCP function calls, then loop with results as next input fn_calls = [o for o in interaction.outputs if o.type == "function_call"] if not fn_calls: # No client-side function calls (e.g. only built-in tool calls) — done break results = [] for fc in fn_calls: print(f"[orchestrator] Calling tool: {fc.name}...", file=sys.stderr) mcp_result = await session.call_tool(fc.name, fc.arguments) result_text = "\n".join( c.text for c in mcp_result.content if hasattr(c, "text") ) or "(no output)" results.append({ "type": "function_result", "call_id": fc.id, "name": fc.name, "result": result_text }) current_input = results elif interaction.status == "failed": raise RuntimeError(f"Interaction failed: {interaction.error}") else: raise RuntimeError(f"Unexpected interaction status: {interaction.status}")
async def invoke_gemini_api( prompt: str, model: str = "gemini-3.1-pro-preview", thinking_level: str = "medium", timeout: int = 900, conversation_id: str | None = None,): """Invoke Gemini via Interactions API with MCP tool bridging.""" start_time = time.time() session_id = f"gemini-{datetime.now().strftime('%Y%m%dT%H%M%S')}" # Reset turn state circuit breaker turn_state_path = HOME / "logs" / ".turn_state.json" turn_state_path.parent.mkdir(parents=True, exist_ok=True) current_state = {} if turn_state_path.exists(): try: current_state = json.loads(turn_state_path.read_text(encoding="utf-8")) except Exception: pass current_state.update({"count": 0, "messages": []}) turn_state_path.write_text(json.dumps(current_state), encoding="utf-8")
prior_size = MESSAGES_LOG.stat().st_size if MESSAGES_LOG.exists() else 0
print(f"[orchestrator] Invoking Gemini API ({model}, {thinking_level}, {len(prompt)} chars)...", file=sys.stderr)
return_code = 0 try: await asyncio.wait_for( _run_gemini_loop(prompt, model, thinking_level, conversation_id), timeout=timeout ) except asyncio.TimeoutError: print(f"[orchestrator] Gemini API timed out after {timeout}s", file=sys.stderr) return_code = 1 except Exception as e: print(f"[orchestrator] Gemini API error: {e}", file=sys.stderr) import traceback traceback.print_exc(file=sys.stderr) return_code = 1
duration = time.time() - start_time log_event("gemini-api", len(prompt), session_id, return_code, duration)
# Capture messages new_messages = read_new_messages(MESSAGES_LOG, prior_size) for msg in new_messages: print(f"Dot: {msg}")
write_session_log("gemini-api", prompt, session_id, new_messages, duration, return_code)
if __name__ == "__main__": # For testing: uv run python invoke_gemini_api.py "hello" if len(sys.argv) > 1: asyncio.run(invoke_gemini_api(sys.argv[1]))