Spaces:
Running
Running
File size: 4,791 Bytes
6a5e36e | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | 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.
|