engine_bitai.py

import os
import asyncio
import re
from typing import TypedDict, List, Annotated, Sequence
from langchain_groq import ChatGroq
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, END
from db import get_service_client
import engine_user
import engine_messaging
import engine_content

# --- State Definition ---
class AgentState(TypedDict):
    messages: List[BaseMessage]
    username: str
    ai_name: str
    personality: str
    context: str
    response: str

# --- RAG Logic ---
def get_user_context(username: str) -> str:
    """Collects recent posts and info about the user to provide context."""
    try:
        # Get user bio and basic info
        user = engine_user.get_user_by_username(username)
        bio = user.get("bio", "No bio provided.")
        
        # Get recent posts by the user
        posts = engine_content.get_user_posts(username, viewer=username)
        recent_posts = "\n".join([f"- {p['content']}" for p in posts[:5]])
        
        context = f"User Bio: {bio}\n\nRecent Posts by {username}:\n{recent_posts}"
        return context
    except Exception as e:
        print(f"Error gathering context: {e}")
        return "No specific context found."

# --- Nodes ---

def retrieve_node(state: AgentState):
    """Retrieve user-specific context."""
    username = state["username"]
    context = get_user_context(username)
    return {"context": context}

def generate_node(state: AgentState):
    """Generate response using context, personality, and history."""
    llm = ChatGroq(
        model="llama-3.1-8b-instant",
        groq_api_key=os.environ.get("GROQ_API_KEY"),
        temperature=0.7,
        max_tokens=256
    )
    
    username = state["username"]
    ai_name = state["ai_name"]
    personality = state["personality"]
    context = state["context"]
    
    system_prompt = f"""You are {ai_name}, a personal AI assistant for {username} on YapStation.
Your Personality/Instructions: {personality if personality else "Be a friendly and helpful cyberpunk assistant."}

YapStation Context:
- A cyberpunk social media platform.
- Features: Stations (podcasts), Posts, Real-time Chat, Stories.

User Context (Recent activity/bio):
{context}

Guidelines:
- If this is the first interaction and they are giving you a name, acknowledge it warmly.
- Always refer to yourself as {ai_name}.
- Keep responses relatively concise but engaging.
- Use cyberpunk slang occasionally if it fits.
"""
    
    messages = [SystemMessage(content=system_prompt)] + state["messages"]
    
    response = llm.invoke(messages)
    return {"response": response.content}

# --- Graph Construction ---

def create_bitai_graph():
    workflow = StateGraph(AgentState)
    
    workflow.add_node("retrieve", retrieve_node)
    workflow.add_node("generate", generate_node)
    
    workflow.set_entry_point("retrieve")
    workflow.add_edge("retrieve", "generate")
    workflow.add_edge("generate", END)
    
    return workflow.compile()

# --- Execution ---

async def run_bitai_chat(username: str, user_message: str, history: List[dict], settings: dict):
    """Runs the LangGraph for a BitAI chat interaction."""
    graph = create_bitai_graph()
    
    # Convert history to LangChain messages
    messages = []
    for m in history:
        if m.get("role") == "user":
            messages.append(HumanMessage(content=m.get("content", "")))
        else:
            messages.append(AIMessage(content=m.get("content", "")))
    
    messages.append(HumanMessage(content=user_message))
    
    inputs = {
        "messages": messages,
        "username": username,
        "ai_name": settings.get("ai_name", "BitAI"),
        "personality": settings.get("personality", ""),
        "context": ""
    }
    
    result = await asyncio.to_thread(graph.invoke, inputs)
    return result["response"]

def extract_name_from_sentence(sentence: str) -> str:
    """Helper to extract a name if the user enters a sentence like 'Your name is JARVIS'."""
    # Simple heuristic: look for "name is [Name]" or "call you [Name]"
    match = re.search(r"name\s+is\s+([A-Za-z0-9\s]+)", sentence, re.IGNORECASE)
    if match:
        return match.group(1).strip().split('\n')[0].split('.')[0]
    
    match = re.search(r"call\s+you\s+([A-Za-z0-9\s]+)", sentence, re.IGNORECASE)
    if match:
        return match.group(1).strip().split('\n')[0].split('.')[0]
    
    # If no pattern, just take the first 2 words if it's short, or the whole thing if it's one word
    words = sentence.strip().split()
    if len(words) <= 2:
        return " ".join(words)
    
    return words[-1] # Fallback to last word? Or just return the whole thing?
    # User said: "I want to call you Jarvis" -> extract Jarvis.