import os import json import asyncio import logging from typing import AsyncGenerator try: import aiohttp except ImportError: logging.error("aiohttp required for LLM Connectors") class AuthError(Exception): pass class RateLimitError(Exception): pass class LLMConnector: def __init__(self, api_key: str): self.api_key = api_key def truncate_context(self, system_prompt: str, prompt: str, max_chars: int = 120000) -> str: # 1 token ~= 4 characters. 120,000 chars ~= 30,000 tokens (very safe limit) total_len = len(system_prompt) + len(prompt) if total_len > max_chars: excess = total_len - max_chars # Truncate the user prompt (which contains history) from the top (oldest memories) # Find the first newline after the excess to avoid cutting mid-sentence idx = prompt.find('\n', excess) if idx == -1: idx = excess prompt = "... [Context Truncated] ...\n" + prompt[idx:] return prompt async def stream_generate(self, system_prompt: str, prompt: str) -> AsyncGenerator[str, None]: raise NotImplementedError class GeminiConnector(LLMConnector): async def stream_generate(self, system_prompt: str, prompt: str) -> AsyncGenerator[str, None]: prompt = self.truncate_context(system_prompt, prompt) url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-3.5-flash:streamGenerateContent?key={self.api_key}" headers = {"Content-Type": "application/json"} payload = { "systemInstruction": {"parts": [{"text": system_prompt}]}, "contents": [{"parts": [{"text": prompt}]}] } for attempt in range(3): async with aiohttp.ClientSession() as session: async with session.post(url, headers=headers, json=payload) as resp: if resp.status == 401 or resp.status == 403: raise AuthError("Invalid Gemini API Key") elif resp.status == 429: logging.warning("Gemini 429 Rate Limit. Backing off 2s...") await asyncio.sleep(2 ** attempt) continue elif resp.status >= 500: raise Exception(f"Gemini Server Error {resp.status}") async for line in resp.content: line = line.decode('utf-8').strip() if line.startswith('"text":'): try: text = line.split('"text":')[1].strip(' ",\n') text = text.replace('\\n', '\n').replace('\\"', '"') yield text except: pass break # Success class OpenAIConnector(LLMConnector): async def stream_generate(self, system_prompt: str, prompt: str) -> AsyncGenerator[str, None]: prompt = self.truncate_context(system_prompt, prompt) url = "https://api.openai.com/v1/chat/completions" headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"} payload = { "model": "gpt-4o", "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt} ], "stream": True } for attempt in range(3): async with aiohttp.ClientSession() as session: async with session.post(url, headers=headers, json=payload) as resp: if resp.status == 401: raise AuthError("Invalid OpenAI API Key") elif resp.status == 429: await asyncio.sleep(2 ** attempt) continue elif resp.status >= 500: raise Exception(f"OpenAI Server Error {resp.status}") async for line in resp.content: line = line.decode('utf-8').strip() if line.startswith("data: ") and line != "data: [DONE]": try: data = json.loads(line[6:]) delta = data['choices'][0]['delta'].get('content', '') if delta: yield delta except: pass break class AnthropicConnector(LLMConnector): async def stream_generate(self, system_prompt: str, prompt: str) -> AsyncGenerator[str, None]: prompt = self.truncate_context(system_prompt, prompt) url = "https://api.anthropic.com/v1/messages" headers = { "x-api-key": self.api_key, "anthropic-version": "2023-06-01", "Content-Type": "application/json" } payload = { "model": "claude-3-5-sonnet-20240620", "system": system_prompt, "messages": [{"role": "user", "content": prompt}], "max_tokens": 1024, "stream": True } for attempt in range(3): async with aiohttp.ClientSession() as session: async with session.post(url, headers=headers, json=payload) as resp: if resp.status == 401: raise AuthError("Invalid Anthropic API Key") elif resp.status == 429: await asyncio.sleep(2 ** attempt) continue elif resp.status >= 500: raise Exception(f"Anthropic Server Error {resp.status}") async for line in resp.content: line = line.decode('utf-8').strip() if line.startswith("data: "): try: data = json.loads(line[6:]) if data.get('type') == 'content_block_delta': yield data['delta']['text'] except: pass break class GroqConnector(OpenAIConnector): async def stream_generate(self, system_prompt: str, prompt: str) -> AsyncGenerator[str, None]: prompt = self.truncate_context(system_prompt, prompt) url = "https://api.groq.com/openai/v1/chat/completions" headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"} payload = { "model": "llama-3.1-70b-versatile", "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt} ], "stream": True } for attempt in range(3): async with aiohttp.ClientSession() as session: async with session.post(url, headers=headers, json=payload) as resp: if resp.status == 401: raise AuthError("Invalid Groq API Key") elif resp.status == 429: await asyncio.sleep(2 ** attempt) continue elif resp.status >= 500: raise Exception(f"Groq Error") async for line in resp.content: line = line.decode('utf-8').strip() if line.startswith("data: ") and line != "data: [DONE]": try: yield json.loads(line[6:])['choices'][0]['delta'].get('content', '') except: pass break class OllamaConnector(LLMConnector): async def stream_generate(self, system_prompt: str, prompt: str) -> AsyncGenerator[str, None]: prompt = self.truncate_context(system_prompt, prompt) url = "http://localhost:11434/api/chat" payload = { "model": "llama3.1", "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt} ] } async with aiohttp.ClientSession() as session: async with session.post(url, json=payload) as resp: if resp.status >= 400: raise Exception(f"Ollama Error") async for line in resp.content: try: yield json.loads(line.decode('utf-8'))['message'].get('content', '') except: pass class VLLMConnector(OpenAIConnector): async def stream_generate(self, system_prompt: str, prompt: str) -> AsyncGenerator[str, None]: prompt = self.truncate_context(system_prompt, prompt) url = "http://localhost:8000/v1/chat/completions" payload = { "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt} ], "stream": True } async with aiohttp.ClientSession() as session: async with session.post(url, json=payload) as resp: if resp.status >= 400: raise Exception("vLLM Error") async for line in resp.content: line = line.decode('utf-8').strip() if line.startswith("data: ") and line != "data: [DONE]": try: yield json.loads(line[6:])['choices'][0]['delta'].get('content', '') except: pass async def stream_with_fallback(system_prompt: str, prompt: str) -> AsyncGenerator[str, None]: # ITEM: Both JARVIS and FRIDAY use Gemini as the primary brain natively import os try: from config import GEMINI_API_KEY except ImportError: GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "") # Priority List: Gemini -> Groq -> OpenAI -> Ollama keys = { "gemini": GEMINI_API_KEY, "groq": os.environ.get("GROQ_API_KEY", ""), "openai": os.environ.get("OPENAI_API_KEY", "") } connectors = [ (GeminiConnector(keys["gemini"]), "Gemini"), (GroqConnector(keys["groq"]), "Groq"), (OpenAIConnector(keys["openai"]), "OpenAI"), (OllamaConnector(""), "Ollama") ] last_error = "" for connector, name in connectors: if isinstance(connector, OllamaConnector) or connector.api_key: logging.info(f"Routing to {name}...") try: stream_active = False async for token in connector.stream_generate(system_prompt, prompt): stream_active = True yield token if stream_active: return # Successfully finished streaming except AuthError as e: # ITEM: 401 Error Recovery - Do not retry, emit error logging.error(f"{name} 401 Auth Error: {e}") yield f"\n\n[SYSTEM: {name} API Key is Invalid. Aborting.]" return except Exception as e: logging.warning(f"{name} Failed: {e}. Falling back to next connector...") last_error = str(e) continue yield f"\n\n[SYSTEM: All Cognitive Connectors Offline. Last Error: {last_error}]"