diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 763 |
1 files changed, 763 insertions, 0 deletions
@@ -0,0 +1,763 @@ +import os +import sys +import json +import subprocess +import hashlib +import re +import threading +import datetime # <-- NEW +from abc import ABC, abstractmethod +from typing import List, Dict, Any, Generator, Optional, Callable +from dotenv import load_dotenv + +import ollama +from cerebras.cloud.sdk import Cerebras +from ddgs import DDGS +from langchain_ollama import OllamaEmbeddings +from langchain_community.vectorstores import FAISS +from langchain_core.documents import Document +from langchain_text_splitters import RecursiveCharacterTextSplitter + +from rich.console import Console +from rich.prompt import Prompt +from rich.panel import Panel + +load_dotenv() + +# ========================================== +# 1. Configuration & Constants +# ========================================== +LOCAL_LLM = "qwen3-vl:8b" +LOCAL_EMBED_MODEL = "nomic-embed-text-v2-moe:latest" +PKM_DIR = os.path.expanduser("~/monorepo") + +XDG_CONFIG_HOME = os.environ.get("XDG_CONFIG_HOME", os.path.expanduser("~/.config")) +APP_CONFIG_DIR = os.path.join(XDG_CONFIG_HOME, "cerebral") +APP_CACHE_DIR = os.path.expanduser("~/.cache/cerebral") +ORG_OUTPUT_DIR = os.path.expanduser("~/org/cerebral") + +os.makedirs(APP_CONFIG_DIR, exist_ok=True) +os.makedirs(APP_CACHE_DIR, exist_ok=True) +os.makedirs(ORG_OUTPUT_DIR, exist_ok=True) + +MEMORY_FILE = os.path.join(APP_CACHE_DIR, "memory_summary.txt") +MEMORY_INDEX_PATH = os.path.join(APP_CACHE_DIR, "memory_index") +FAISS_INDEX_PATH = os.path.join(APP_CONFIG_DIR, "pkm_index") +HASH_TRACKER_FILE = os.path.join(APP_CONFIG_DIR, "latest_commit.txt") + +# ========================================== +# 2. Abstract LLM Provider +# ========================================== +class BaseLLMProvider(ABC): + """Abstract interface for LLM providers to ensure easy swapping.""" + @abstractmethod + # <-- UPDATED: Added tool_choice parameter + def chat_completion(self, messages: List[Dict], tools: List[Dict] = None, stream: bool = False, tool_choice: str = "auto") -> Any: + pass + +class CerebrasProvider(BaseLLMProvider): + def __init__(self, model: str = "qwen-3-235b-a22b-instruct-2507"): + api_key = os.environ.get("CEREBRAS_API_KEY") + if not api_key: + raise ValueError("CEREBRAS_API_KEY environment variable is required.") + self.client = Cerebras(api_key=api_key) + self.model = model + + def chat_completion(self, messages: List[Dict], tools: List[Dict] = None, stream: bool = False, tool_choice: str = "auto"): + kwargs = { + "messages": messages, + "model": self.model, + "stream": stream, + } + if tools: + kwargs["tools"] = tools + kwargs["tool_choice"] = tool_choice # <-- UPDATED + + return self.client.chat.completions.create(**kwargs) + +# ========================================== +# 3. Core Modules +# ========================================== +class MemoryManager: + def __init__(self, memory_file: str, index_path: str, local_model: str, embed_model_name: str, log: Callable[[str], None] = print): + self.memory_file = memory_file + self.index_path = index_path + self.local_model = local_model + self.log = log + + self.session_summary = "Session just started. No prior context." + self.interaction_buffer = [] + self.COMPRESSION_THRESHOLD = 4 + + self.embeddings = OllamaEmbeddings(model=embed_model_name) + self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50) + + self.log("[dim italic]Loading persistent memory...[/dim italic]") + if os.path.exists(self.memory_file): + with open(self.memory_file, "r") as f: + self.persistent_memory = f.read().strip() + else: + self.persistent_memory = "No known user facts or long-term preferences." + + if os.path.exists(self.index_path): + self.vectorstore = FAISS.load_local(self.index_path, self.embeddings, allow_dangerous_deserialization=True) + else: + self.log("[bold yellow]No memory index found. Building initial database...[/bold yellow]") + self.rebuild_index() + + def get_line_count(self) -> int: + if not os.path.exists(self.memory_file): + return 0 + with open(self.memory_file, "r") as f: + return sum(1 for _ in f) + + def rebuild_index(self): + self.log("[dim italic]Reserializing memory log into vector database...[/dim italic]") + text = self.persistent_memory if self.persistent_memory else "No known user facts or long-term preferences." + + chunks = self.text_splitter.split_text(text) + docs = [Document(page_content=c) for c in chunks] + self.vectorstore = FAISS.from_documents(docs, self.embeddings) + self.vectorstore.save_local(self.index_path) + self.log("[bold green]Memory database manually rebuilt and saved![/bold green]") + + def compress_persistent_memory(self): + self.log("[bold yellow]Compressing persistent memory (removing duplicates and irrelevant data)...[/bold yellow]") + if not os.path.exists(self.memory_file): + self.log("[dim]Memory file is empty. Nothing to compress.[/dim]") + return + + # STRICT PROMPT FOR COMPRESSION + sys_prompt = """You are a strictly robotic data deduplication script. Your ONLY job is to compress the provided memory log. + RULES: + 1. Remove duplicate facts. + 2. Remove conversational text, essays, or philosophical analysis. + 3. Output ONLY a clean, simple bulleted list of facts. + 4. NEVER use headers, bold text, or introductory/closing remarks.""" + + try: + response = ollama.chat(model=self.local_model, messages=[ + {'role': 'system', 'content': sys_prompt}, + {'role': 'user', 'content': f"MEMORY LOG TO COMPRESS:\n{self.persistent_memory}"} + ]) + compressed_memory = response['message']['content'].strip() + compressed_memory = re.sub(r'<think>.*?</think>', '', compressed_memory, flags=re.DOTALL).strip() + + with open(self.memory_file, "w") as f: + f.write(compressed_memory) + self.persistent_memory = compressed_memory + self.rebuild_index() + self.log("[bold green]Persistent memory successfully compressed and re-indexed![/bold green]") + except Exception as e: + self.log(f"[bold red]Memory compression failed: {e}[/bold red]") + + def search(self, query: str) -> str: + if not getattr(self, 'vectorstore', None): + return "No long-term memories available." + docs = self.vectorstore.similarity_search(query, k=3) + return "\n".join([f"- {d.page_content}" for d in docs]) + + def add_interaction(self, user_input: str, bot_response: str): + self.interaction_buffer.append({"user": user_input, "agent": bot_response}) + if len(self.interaction_buffer) >= self.COMPRESSION_THRESHOLD: + buffer_to_compress = list(self.interaction_buffer) + self.interaction_buffer = [] + threading.Thread(target=self._compress_session, args=(buffer_to_compress,), daemon=True).start() + + def _compress_session(self, buffer: List[Dict]): + buffer_text = "\n".join([f"User: {i['user']}\nAgent: {i['agent']}" for i in buffer]) + + # STRICT PROMPT FOR SESSION COMPRESSION + sys_prompt = """You are a strict summarization script. Merge the recent interactions into the current session summary. + RULES: + 1. Keep it brief and objective. + 2. DO NOT write essays or analyze the user's intent. + 3. Output ONLY the raw text of the updated summary. No conversational padding.""" + + try: + response = ollama.chat(model=self.local_model, messages=[ + {'role': 'system', 'content': sys_prompt}, + {'role': 'user', 'content': f"CURRENT SUMMARY:\n{self.session_summary}\n\nNEW INTERACTIONS:\n{buffer_text}"} + ]) + self.session_summary = response['message']['content'].strip() + self.session_summary = re.sub(r'<think>.*?</think>', '', self.session_summary, flags=re.DOTALL).strip() + except Exception as e: + self.log(f"[dim red]Background session compression failed: {e}[/dim red]") + + def finalize_session(self): + self.log("[bold yellow]Extracting long-term memories from session...[/bold yellow]") + final_context = self.session_summary + if self.interaction_buffer: + final_context += "\n" + "\n".join([f"User: {i['user']}\nAgent: {i['agent']}" for i in self.interaction_buffer]) + + # STRICT PROMPT FOR EXTRACTION + sys_prompt = """You are a strict data extraction pipeline. Your ONLY job is to extract permanent, long-term facts about the user from the provided session text. + + RULES: + 1. NEVER write conversational text, greetings, headers, or explanations. + 2. NEVER write essays, evaluate, or analyze the meaning of the facts. + 3. ONLY output a raw, bulleted list of concise facts (e.g., "- User uses Emacs org-mode"). + 4. If there are NO new permanent facts to save, output EXACTLY and ONLY the word: NONE. + """ + + try: + response = ollama.chat(model=self.local_model, messages=[ + {'role': 'system', 'content': sys_prompt}, + {'role': 'user', 'content': f"SESSION TEXT TO EXTRACT FROM:\n{final_context}"} + ]) + + new_facts = response['message']['content'].strip() + new_facts = re.sub(r'<think>.*?</think>', '', new_facts, flags=re.DOTALL).strip() + + if new_facts.upper() != "NONE" and new_facts: + # Failsafe: If the model hallucinates an essay anyway, block it from saving. + if len(new_facts.split('\n')) > 15 or "###" in new_facts: + self.log("[dim red]Model hallucinated an essay instead of facts. Discarding to protect memory database.[/dim red]") + return + + with open(self.memory_file, "a") as f: + f.write(f"\n{new_facts}") + self.persistent_memory += f"\n{new_facts}" + self.log("[bold green]New facts appended to long-term memory log![/bold green]") + self.log("[dim]Note: Run /memory rebuild to index these new facts for next time.[/dim]") + else: + self.log("[dim]No new long-term facts detected. Skipping memory append.[/dim]") + except Exception as e: + self.log(f"[bold red]Failed to save long-term memory: {e}[/bold red]") + +class PKMManager: + def __init__(self, pkm_dir: str, index_path: str, hash_file: str, embed_model_name: str, log: Callable[[str], None] = print): + self.pkm_dir = pkm_dir + self.index_path = index_path + self.hash_file = hash_file + self.log = log + + self.log(f"[dim italic]Waking up Ollama embeddings ({embed_model_name})...[/dim italic]") + self.embeddings = OllamaEmbeddings(model=embed_model_name) + self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + self.vectorstore = self._load_or_build() + + def _get_main_commit_hash(self) -> str: + try: + result = subprocess.run( + ["git", "rev-parse", "main"], + cwd=self.pkm_dir, capture_output=True, text=True, check=True + ) + return result.stdout.strip() + except subprocess.CalledProcessError: + return "unknown" + + def _load_or_build(self): + self.log("[dim]Checking Git HEAD hash for PKM changes...[/dim]") + current_hash = self._get_main_commit_hash() + + if os.path.exists(self.index_path) and os.path.exists(self.hash_file): + with open(self.hash_file, "r") as f: + if f.read().strip() == current_hash: + self.log(f"[green]Git hash unchanged ({current_hash[:7]}). Loading cached PKM index...[/green]") + return FAISS.load_local(self.index_path, self.embeddings, allow_dangerous_deserialization=True) + + self.log(f"[bold yellow]New commits detected ({current_hash[:7]}). Rebuilding PKM index...[/bold yellow]") + raw_documents = [] + + self.log(f"[dim]Scanning {self.pkm_dir} for .org files...[/dim]") + for root, dirs, files in os.walk(self.pkm_dir): + if '.git' in dirs: dirs.remove('.git') + if 'nix' in dirs: dirs.remove('nix') + for file in files: + if file.endswith('.org'): + filepath = os.path.join(root, file) + try: + with open(filepath, 'r', encoding='utf-8') as f: + raw_documents.append(Document(page_content=f.read(), metadata={"source": filepath})) + except Exception: + pass + + if not raw_documents: + self.log("[red]No .org files found in PKM directory.[/red]") + return None + + self.log(f"[dim]Chunking {len(raw_documents)} documents...[/dim]") + chunks = self.text_splitter.split_documents(raw_documents) + + self.log(f"[bold cyan]Embedding {len(chunks)} chunks via Ollama (this might take a minute)...[/bold cyan]") + vectorstore = FAISS.from_documents(chunks, self.embeddings) + vectorstore.save_local(self.index_path) + + with open(self.hash_file, "w") as f: + f.write(current_hash) + + self.log("[bold green]PKM Index successfully rebuilt and saved![/bold green]") + return vectorstore + + def search(self, query: str) -> str: + if not self.vectorstore: + return "PKM is empty." + docs = self.vectorstore.similarity_search(query, k=10) + return "PKM Search Results:\n" + "\n\n".join([f"From {d.metadata['source']}:\n{d.page_content}" for d in docs]) + + +class VisionProcessor: + def __init__(self, local_model: str, log: Callable[[str], None] = print): + self.local_model = local_model + self.log = log + self.log("[dim italic]Vision Processor online...[/dim italic]") + + def process(self, image_path: str, user_prompt: str) -> str: + try: + with open(image_path, 'rb') as img_file: + img_bytes = img_file.read() + response = ollama.chat(model=self.local_model, messages=[{ + 'role': 'user', + 'content': f"Describe this image in detail to help another AI answer this prompt: {user_prompt}", + 'images': [img_bytes] + }]) + return response['message']['content'] + except Exception as e: + return f"[Image analysis failed: {e}]" + + +# ========================================== +# Web Search Providers +# ========================================== +class BaseSearchProvider(ABC): + """Abstract interface for web search engines to ensure easy swapping and fallbacks.""" + @abstractmethod + def search(self, query: str, max_results: int = 10) -> List[Dict[str, str]]: + pass + +class GoogleSearchProvider(BaseSearchProvider): + def search(self, query: str, max_results: int = 10) -> List[Dict[str, str]]: + # Imported locally so it doesn't crash the app if the package is missing + from googlesearch import search + + results = [] + # advanced=True forces it to return objects with title, url, and description + for r in search(query, num_results=max_results, advanced=True): + results.append({ + 'title': getattr(r, 'title', 'No Title'), + 'href': getattr(r, 'url', 'No URL'), + 'body': getattr(r, 'description', 'No Description') + }) + + if not results: + raise Exception("Google returned zero results.") + + return results + +class DDGSSearchProvider(BaseSearchProvider): + def search(self, query: str, max_results: int = 10) -> List[Dict[str, str]]: + results = DDGS().text(query, max_results=max_results) + + if not results: + raise Exception("DuckDuckGo returned zero results.") + + formatted_results = [] + for r in results: + formatted_results.append({ + 'title': r.get('title', 'No Title'), + 'href': r.get('href', 'No URL'), + 'body': r.get('body', 'No Description') + }) + return formatted_results + +class WebSearcher: + def __init__(self, log: Callable[[str], None] = print): + self.log = log + # The order of this list dictates the fallback priority + self.providers: List[BaseSearchProvider] = [ + GoogleSearchProvider(), + DDGSSearchProvider() + ] + + def search(self, query: str) -> str: + for provider in self.providers: + provider_name = provider.__class__.__name__ + try: + self.log(f"[dim italic]Trying {provider_name}...[/dim italic]") + results = provider.search(query, max_results=10) + + context = "Web Search Results:\n" + for r in results: + context += f"- Title: {r['title']}\n URL: {r['href']}\n Snippet: {r['body']}\n\n" + return context + + except Exception as e: + # Catch 429 Rate Limits, connection errors, or empty results and seamlessly fall back + self.log(f"[dim yellow]{provider_name} failed ({e}). Falling back...[/dim yellow]") + continue + + return "Web search failed: All search providers were exhausted or rate-limited." + +# ========================================== +# 4. The Orchestrator (Agnostic Agent) +# ========================================== +class CerebralAgent: + def __init__(self, provider: BaseLLMProvider, log: Callable[[str], None] = print): + self.provider = provider + self.log = log + + self.log("[bold magenta]Initializing Cerebral Agent Modules...[/bold magenta]") + self.memory = MemoryManager(MEMORY_FILE, MEMORY_INDEX_PATH, LOCAL_LLM, LOCAL_EMBED_MODEL, self.log) + self.pkm = PKMManager(PKM_DIR, FAISS_INDEX_PATH, HASH_TRACKER_FILE, LOCAL_EMBED_MODEL, self.log) + self.vision = VisionProcessor(LOCAL_LLM, self.log) + self.web = WebSearcher(self.log) + + def generate_session_filename(self, first_prompt: str, first_response: str) -> str: + self.log("[dim italic]Generating descriptive filename based on prompt and response...[/dim italic]") + + hash_input = (first_prompt + first_response).encode('utf-8') + combined_hash = hashlib.sha256(hash_input).hexdigest()[:6] + + sys_prompt = "You are a file naming utility. Read the user's prompt and generate a short, descriptive filename base using ONLY lowercase letters and hyphens. Do NOT add an extension. ONLY output the base filename, absolutely no other text. Example: learning-python-basics" + + try: + response = ollama.chat(model=LOCAL_LLM, messages=[ + {'role': 'system', 'content': sys_prompt}, + {'role': 'user', 'content': first_prompt} + ]) + + raw_content = response['message']['content'].strip() + raw_content = re.sub(r'<think>.*?</think>', '', raw_content, flags=re.DOTALL).strip() + + lines = [line.strip() for line in raw_content.split('\n') if line.strip()] + raw_filename = lines[-1].lower().replace(' ', '-') if lines else "cerebral-session" + clean_base = re.sub(r'[^a-z0-9\-]', '', raw_filename).strip('-') + clean_base = clean_base[:50].strip('-') + + if not clean_base: + clean_base = "cerebral-session" + + final_filename = f"{clean_base}-{combined_hash}.org" + return final_filename + + except Exception as e: + self.log(f"[dim red]Filename generation failed: {e}. Defaulting.[/dim red]") + return f"cerebral-session-{combined_hash}.org" + + def _get_tools(self) -> List[Dict]: + return [ + { + "type": "function", + "function": { + "name": "search_pkm", + "description": "Search the user's personal knowledge base (PKM) for notes, code, or org files.", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]} + } + }, + { + "type": "function", + "function": { + "name": "search_web", + "description": "Search the live internet for current events, external documentation, or facts outside your PKM.", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]} + } + } + ] + + def chat_stream(self, prompt: str, image_path: Optional[str] = None) -> Generator[str, None, str]: + """Core interaction loop. Yields text chunks. Returns full text when done.""" + + recent_history = "" + if self.memory.interaction_buffer: + recent_history = "\nRECENT UNCOMPRESSED TURNS:\n" + "\n".join( + [f"User: {i['user']}\nAgent: {i['agent']}" for i in self.memory.interaction_buffer] + ) + + vision_context = "" + if image_path: + self.log("[dim italic]Analyzing image context...[/dim italic]") + vision_summary = self.vision.process(image_path, prompt) + vision_context = f"\n[USER ATTACHED AN IMAGE. Local Vision Summary: {vision_summary}]\n" + + self.log("[dim italic]Querying long-term memory (Ollama Embeddings)...[/dim italic]") + relevant_memories = self.memory.search(prompt) + + current_time = datetime.datetime.now().strftime("%A, %B %d, %Y at %I:%M %p") + + system_prompt = f"""You are a highly capable AI assistant. + + CRITICAL OUTPUT FORMATTING: + You MUST output your responses EXCLUSIVELY in Emacs org-mode format. Use org-mode headings, lists, and LaTeX fragments for math. + + FORMATTING RULES: + 1. NEVER use double asterisks (`**`) for bolding. You MUST use SINGLE asterisks for bold emphasis (e.g., *this is bold*). Double asterisks will break the parser. + 2. Cite your sources inline using proper org-mode link syntax. For web searches, use [[url][Description]]. For PKM files, use [[file:/path/to/file.org][Filename]]. + 3. At the very end of your response, you MUST append a Level 1 heading `* Sources` and neatly list all the search results and PKM documents you referenced using proper org-mode syntax. + + CURRENT TIME AND DATE: + {current_time} + + RESPONSE STYLE GUIDELINES: + - Provide EXTREMELY detailed, exhaustive, and comprehensive answers. + - Write in long-form prose. Do not be brief; expand deeply on concepts. + - Use multiple paragraphs, deep conceptual explanations, and thorough analysis. + + RELEVANT LONG-TERM MEMORIES: + {relevant_memories} + + COMPRESSED SESSION CONTEXT: {self.memory.session_summary} + {recent_history} + """ + + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt + vision_context} + ] + + self.log("[dim italic]Analyzing intent & tool requirements (Cerebras)...[/dim italic]") + + # --- NEW: Self-Healing Tool Call Loop --- + MAX_RETRIES = 3 + valid_tool_calls = False + response_message = None + allowed_tool_names = [t["function"]["name"] for t in self._get_tools()] + + for attempt in range(MAX_RETRIES): + pre_flight = self.provider.chat_completion(messages=messages, tools=self._get_tools(), stream=False) + response_message = pre_flight.choices[0].message + + # Scenario A: Hallucinated Markdown Tool Call + if not response_message.tool_calls and response_message.content and "**name**:" in response_message.content: + self.log(f"[dim yellow]Model hallucinated Markdown tool call. Retrying ({attempt+1}/{MAX_RETRIES})...[/dim yellow]") + error_msg = f"ERROR: You attempted to call a tool using Markdown text. You MUST use the native JSON tool calling API. Allowed tools: {allowed_tool_names}" + messages.append({"role": "assistant", "content": response_message.content}) + messages.append({"role": "user", "content": error_msg}) + continue + + # Scenario B: Legitimate text response (No tools needed) + if not response_message.tool_calls: + valid_tool_calls = True + break + + # Scenario C: Native API Tool Calls (Needs Validation) + has_errors = False + error_feedbacks = [] + + for tool_call in response_message.tool_calls: + func_name = tool_call.function.name + call_error = None + + if func_name not in allowed_tool_names: + has_errors = True + call_error = f"Tool '{func_name}' does not exist. Allowed tools: {allowed_tool_names}" + else: + try: + json.loads(tool_call.function.arguments) + except json.JSONDecodeError: + has_errors = True + call_error = f"Arguments for '{func_name}' are not valid JSON: {tool_call.function.arguments}" + + error_feedbacks.append(call_error) + + if has_errors: + self.log(f"[dim yellow]Malformed tool call detected. Retrying ({attempt+1}/{MAX_RETRIES})...[/dim yellow]") + + # Append the bad tool call to history so it learns what it did wrong + assistant_msg = { + "role": "assistant", + "content": response_message.content or "", + "tool_calls": [ + { + "id": t.id, + "type": "function", + "function": {"name": t.function.name, "arguments": t.function.arguments} + } for t in response_message.tool_calls + ] + } + messages.append(assistant_msg) + + # Append the specific errors as API tool responses + for i, tool_call in enumerate(response_message.tool_calls): + err = error_feedbacks[i] + msg_content = f"ERROR: {err}" if err else "Error: Another tool in this batch failed. Please fix the batch and retry." + messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": msg_content}) + continue + + # Scenario D: Valid Tool Calls + valid_tool_calls = True + break + + # Failsafe: If it fails 3 times, wipe the tool calls to force a graceful text degradation + if not valid_tool_calls: + self.log("[bold red]Failed to generate valid tool calls. Proceeding without tools.[/bold red]") + response_message.tool_calls = None + # ---------------------------------------- + + if not response_message.tool_calls: + self.log("[dim italic]No tools needed. Outputting response...[/dim italic]") + content = response_message.content or "" + yield content + self.memory.add_interaction(prompt, content) + return content + + # --- Execute Validated Tools --- + assistant_msg = { + "role": "assistant", + "content": response_message.content or "", + "tool_calls": [ + { + "id": t.id, + "type": "function", + "function": {"name": t.function.name, "arguments": t.function.arguments} + } for t in response_message.tool_calls + ] + } + messages.append(assistant_msg) + + for tool_call in response_message.tool_calls: + func_name = tool_call.function.name + args = json.loads(tool_call.function.arguments) # Guaranteed to be safe now + + if func_name == "search_pkm": + q = args.get("query", prompt) + self.log(f"[cyan]🧠 Tool Call: Searching PKM for '{q}'...[/cyan]") + yield f"\n*(Agent Note: Searched PKM for `{q}`)*\n\n" + result = self.pkm.search(q) + elif func_name == "search_web": + q = args.get("query", prompt) + self.log(f"[cyan]🌐 Tool Call: Searching Web for '{q}'...[/cyan]") + yield f"\n*(Agent Note: Searched Web for `{q}`)*\n\n" + result = self.web.search(q) + + messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": result}) + + messages.append({ + "role": "system", + "content": "Tool results received. Now provide your final, comprehensive answer in strict org-mode. REMEMBER: Use *single asterisks* for bold, NEVER double asterisks." + }) + + self.log("[dim italic]Streaming final response...[/dim italic]") + stream = self.provider.chat_completion(messages=messages, tools=self._get_tools(), stream=True, tool_choice="none") + + full_response = "" + for chunk in stream: + content = chunk.choices[0].delta.content or "" + full_response += content + yield content + + self.memory.add_interaction(prompt, full_response) + return full_response + + def shutdown(self): + self.memory.finalize_session() + + +# ========================================== +# 5. The CLI Presentation Layer +# ========================================== +class CLIApp: + def __init__(self, agent: CerebralAgent, console: Console): + self.agent = agent + self.console = console + self.current_session_file = None + + def run(self): + self.console.print(Panel.fit("🤖 [bold blue]Modular Cerebral Agent[/bold blue] initialized.\n- Type [bold]/image /path/to/img.png <prompt>[/bold] to attach images.\n- Type [bold]/exit[/bold] to quit.", border_style="blue")) + + while True: + try: + user_input = Prompt.ask("\n[bold magenta]You[/bold magenta]") + + if user_input.lower() == '/memory count': + count = self.agent.memory.get_line_count() + self.console.print(f"[bold cyan]Persistent Memory Lines:[/bold cyan] {count}") + continue + + if user_input.lower() == '/memory rebuild': + self.agent.memory.rebuild_index() + continue + + if user_input.lower() == '/memory compress': + self.agent.memory.compress_persistent_memory() + continue + + if clean_input == '/memory': + help_text = ( + "[bold cyan]/memory count[/bold cyan] : Print the number of lines in persistent memory.\n" + "[bold cyan]/memory rebuild[/bold cyan] : Manually reserialize the FAISS database from the log.\n" + "[bold cyan]/memory compress[/bold cyan] : Use the local LLM to scrub duplicates and compress the log." + ) + self.console.print(Panel.fit(help_text, title="🧠 Memory Commands", border_style="cyan")) + continue + + if user_input.lower() in ['/exit', '/quit']: + self.console.print("\n[dim italic]Initiating shutdown sequence...[/dim italic]") + self.agent.shutdown() + self.console.print("[bold red]Exiting...[/bold red]") + break + + if not user_input.strip(): + continue + + image_path = None + prompt = user_input + if user_input.startswith("/image "): + parts = user_input.split(" ", 2) + if len(parts) >= 2: + image_path = parts[1] + prompt = parts[2] if len(parts) > 2 else "What is this?" + self.console.print(f"[dim italic]Processing image locally...[/dim italic]") + + self.console.print("[bold green]Agent:[/bold green]") + + if not self.current_session_file: + full_response = "" + for chunk in self.agent.chat_stream(prompt, image_path=image_path): + print(chunk, end="", flush=True) + full_response += chunk + print("\n") + + generated_name = self.agent.generate_session_filename(prompt, full_response) + self.current_session_file = os.path.join(ORG_OUTPUT_DIR, generated_name) + + self.console.print(f"[bold green]Session log created at:[/bold green] [cyan]{self.current_session_file}[/cyan]") + + with open(self.current_session_file, "w") as f: + f.write(f"* User Prompt: {user_input}\n** Response\n{full_response}\n") + + try: + self.console.print("[dim italic]Triggering emacsclient...[/dim italic]") + subprocess.run( + ["emacsclient", "-n", self.current_session_file], + check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL + ) + except Exception as e: + self.console.print(f"[dim red]Failed to trigger emacsclient: {e}[/dim red]") + + else: + full_response = "" + with open(self.current_session_file, "a") as f: + f.write(f"\n* User Prompt: {user_input}\n** Response\n") + + for chunk in self.agent.chat_stream(prompt, image_path=image_path): + print(chunk, end="", flush=True) + f.write(chunk) + full_response += chunk + + f.write("\n") + print("\n") + + except KeyboardInterrupt: + self.console.print("\n[bold red]Interrupted. Saving memories...[/bold red]") + self.agent.shutdown() + break + except Exception as e: + self.console.print(f"[bold red]An error occurred: {e}[/bold red]") + +# ========================================== +# 6. Entry Point +# ========================================== +if __name__ == "__main__": + console = Console() + try: + provider = CerebrasProvider() + except ValueError as e: + console.print(f"[bold red]Configuration Error: {e}[/bold red]") + sys.exit(1) + + with console.status("[bold green]Booting up systems...[/bold green]", spinner="dots") as status: + agent = CerebralAgent(provider=provider, log=console.print) + + app = CLIApp(agent, console) + app.run() |
