File size: 4,578 Bytes
6e02dfb | 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 | import os
import json
from google import genai
from openai import OpenAI
class Brain:
def __init__(self):
# --- CONNECTIONS ---
self.gemini_key = os.getenv("GEMINI_API_KEY")
self.openrouter_key = os.getenv("OPENROUTER_API_KEY")
self.gemini_client = None
if self.gemini_key:
self.gemini_client = genai.Client(api_key=self.gemini_key)
self.openrouter_client = None
if self.openrouter_key:
self.openrouter_client = OpenAI(
api_key=self.openrouter_key,
base_url="https://openrouter.ai/api/v1"
)
# --- TIG: THE TOOL REGISTRY ---
self.tool_registry = {
"coding": "anthropic/claude-sonnet-4",
"creative_text": "openai/gpt-5.1-chat",
"research": "gemini-3-flash-preview",
"chat": "gemini-2.5-flash"
}
# State
self.active_model = self.tool_registry["chat"]
self.last_intent = "startup"
def detect_intent(self, user_prompt):
"""
Layer 1: Uses a fast, cheap model to classify the user's intent.
Returns: 'coding', 'creative_text', 'research', or 'chat'
"""
if not self.gemini_client:
return "chat" # Fallback if Google is dead
try:
# Quick classification prompt
classifier_prompt = f"""
ANALYZE this user prompt and output ONLY ONE word from this list:
[coding, creative_text, research, chat]
- coding: asking for python, scripts, html, debugging, logic.
- creative_text: writing stories, poems, complex essays.
- research: asking for facts, summaries of files, history.
- chat: casual conversation, greetings, simple questions.
PROMPT: "{user_prompt[:500]}"
"""
response = self.gemini_client.models.generate_content(
model="gemini-2.5-flash",
contents=classifier_prompt
)
intent = response.text.strip().lower()
# Cleaning result just in case
for valid in self.tool_registry.keys():
if valid in intent:
return valid
return "chat"
except Exception as e:
print(f" [TIG Router] Intent detection failed: {e}")
return "chat"
def think(self, prompt, force_model=None):
# 1. Did user force a model?
target_model = force_model
# 2. If not, Run TIG Intent Detection
if not target_model:
intent = self.detect_intent(prompt)
# --- THE FIX: ALWAYS RESET TO REGISTRY DEFAULT ---
# Previously, we kept 'self.active_model' stuck on the last choice.
# Now, we look up the intent fresh every time.
target_model = self.tool_registry.get(intent, "gemini-2.5-flash")
self.last_intent = intent
self.active_model = target_model # Update state
# Only announce if it's NOT the default chat
if intent != "chat":
print(f" [TIG] Intent: {intent.upper()} -> Routing to {target_model}")
# 3. EXECUTION (The "Do It" Phase)
# PATH A: Use Google Direct (If the chosen model is Gemini)
if "gemini" in target_model and self.gemini_client:
try:
response = self.gemini_client.models.generate_content(
model="gemini-2.5-flash",
contents=prompt
)
return response.text
except Exception as e:
print(f" [Brain] Google failed ({e}). Failover to OpenRouter.")
# Fall through to OpenRouter
# PATH B: Use OpenRouter (For Claude, GPT, or Google Failover)
if self.openrouter_client:
try:
# print(f" [Brain] Calling {target_model} via OpenRouter...")
response = self.openrouter_client.chat.completions.create(
model=target_model,
messages=[{"role": "user", "content": prompt}],
temperature=0.7
)
return response.choices[0].message.content
except Exception as e:
return f"[System Critical] OpenRouter call failed: {e}"
return "[System Critical] No API keys active. Check .env." |