Spaces:
Sleeping
Sleeping
File size: 4,197 Bytes
1d32142 | 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 | import os
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.store.postgres.aio import AsyncPostgresStore
from psycopg_pool import AsyncConnectionPool
from langchain_ollama import ChatOllama
from .state import State
from .router import router
from agents.therapist import TherapistAgent
from agents.logical import LogicalAgent
from agents.classifier import create_classifier
from agents.charity_search import CharitySearchAgent
from agents.agentic_rag import AgenticRAGAgent
from encoders.sealion import SeaLionEncoder
from recommender.vector_store import DonorVectorStore
def create_connection_string() -> str:
"""Build PostgreSQL connection string from environment variables."""
db_host = os.getenv("SUPABASE_DB_HOST", "localhost")
db_port = os.getenv("SUPABASE_DB_PORT", "6543")
db_name = os.getenv("SUPABASE_DB_NAME", "postgres")
db_user = os.getenv("SUPABASE_DB_USER", "postgres")
db_password = os.getenv("SUPABASE_DB_PASSWORD", "")
db_sslmode = os.getenv("SUPABASE_DB_SSLMODE", "require")
return (
f"postgres://{db_user}:{db_password}"
f"@{db_host}:{db_port}/{db_name}"
f"?sslmode={db_sslmode}"
)
def create_async_pool() -> AsyncConnectionPool:
"""Create AsyncConnectionPool with proper settings."""
return AsyncConnectionPool(
conninfo=create_connection_string(),
max_size=20,
kwargs={
"autocommit": True,
"prepare_threshold": None,
}
)
async def build_graph_with_memory():
"""Build the graph with Supabase-backed checkpointer and store."""
# Create async connection pool
pool = create_async_pool()
await pool.open()
# Create checkpointer and store from the pool
checkpointer = AsyncPostgresSaver(pool)
store = AsyncPostgresStore(pool)
# Setup tables for store and checkpointer
print("\n[Setup] Setting up LangGraph store and checkpointer tables...")
await checkpointer.setup()
await store.setup()
print("[OK] Store and checkpointer tables created!\n")
# Use Ollama cloud with API key authentication
api_key = os.getenv('OLLAMA_API_KEY')
if api_key:
llm = ChatOllama(
model="gpt-oss:120b",
base_url="https://ollama.com",
client_kwargs={
"headers": {"Authorization": f"Bearer {api_key}"}
}
)
else:
# Fallback to local Ollama if no API key
llm = ChatOllama(model="gpt-oss:120b-cloud")
# Initialize encoder and vector store for Agentic RAG
encoder = None
vector_store = None
try:
sealion_endpoint = os.getenv("SEALION_ENDPOINT")
if sealion_endpoint:
encoder = SeaLionEncoder(endpoint_url=sealion_endpoint)
vector_store = DonorVectorStore(pool)
print("[OK] Agentic RAG initialized with SeaLion encoder\n")
except Exception as e:
print(f"[WARN] Agentic RAG not available: {e}\n")
# Create Agentic RAG agent
agentic_rag_agent = AgenticRAGAgent(llm, encoder, vector_store)
# Build the graph
graph_builder = StateGraph(State)
graph_builder.add_node("classifier", create_classifier(llm))
graph_builder.add_node("therapist", TherapistAgent(llm))
graph_builder.add_node("logical", LogicalAgent(llm))
graph_builder.add_node("charity_search", CharitySearchAgent(llm))
graph_builder.add_node("agentic_rag", agentic_rag_agent)
graph_builder.add_edge(START, "classifier")
graph_builder.add_conditional_edges(
"classifier",
router,
{
"therapist": "therapist",
"logical": "logical",
"charity_search": "charity_search",
"agentic_rag": "agentic_rag"
}
)
graph_builder.add_edge("therapist", END)
graph_builder.add_edge("logical", END)
graph_builder.add_edge("charity_search", END)
graph_builder.add_edge("agentic_rag", END)
# Compile with store and checkpointer
graph = graph_builder.compile(
checkpointer=checkpointer,
store=store,
)
return graph, store, checkpointer
|