diff --git a/.gitignore b/.gitignore index e0d6a75..6247443 100644 --- a/.gitignore +++ b/.gitignore @@ -153,4 +153,5 @@ cython_debug/ # sqlite database *.db -.env \ No newline at end of file +.env +private.txt \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..bbcf39e --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,94 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project Overview + +A Python multi-agent orchestration system built on LangChain/LangGraph, exposed via three interfaces (REST API, CLI, MCP server) that all share the same agent/routing/DB layer. Observability is via Langfuse; models are accessed through LiteLLM (currently Gemini 2.5 Flash). + +## Commands + +```bash +# Setup +uv venv venv +source venv/bin/activate +touch main.db # create empty sqlite db +uv sync # install deps from pyproject.toml + +# Run REST API (http://127.0.0.1:8000, docs at /docs) +uvicorn app.main:app --reload + +# Run interactive CLI +python -m cli + +# Run MCP server (stdio, for Claude Desktop / Cursor) +./venv/bin/python mcp_server.py + +# Run MCP server (SSE/HTTP, http://localhost:8000/sse) +python -m mcp_server sse + +# Run all tests +PYTHONPATH=. pytest + +# Run a single test file / test +PYTHONPATH=. pytest tests/test_weather_agent.py +PYTHONPATH=. pytest tests/test_weather_agent.py::test_choose_agent_routing +``` + +Use `./venv/bin/python` / `./venv/bin/pytest` explicitly when the venv isn't activated (per `.agents/AGENTS.md`). + +Tests set required env vars (mock API keys, `DATABASE_URL=sqlite:///main_test.db`) in `tests/conftest.py`, and use an in-memory SQLite engine + `app.dependency_overrides` per test module — no `.env` needed to run the suite. + +## Architecture + +### Three entrypoints, one core + +`app/main.py` (FastAPI), `cli.py`, and `mcp_server.py` are thin wrappers around the same core: `app.agents.choose_agent` for routing + LangGraph agent graphs (`car_hire_agent`, `weather_agent`, `generate_image_agent`, `marketing_agent`, `rag_ingest_agent`, `rag_query_agent`) for execution. All three entrypoints duplicate the same per-agent "build initial state / invoke / handle `next_question`" dispatch logic — when changing an agent's state shape or adding an agent, update all three call sites plus `app/agents/__init__.py`. + +### Request flow (REST, mirrored by CLI/MCP) + +1. `POST /llm-job` (`app/routes/llm_job.py`) creates an `LLMJob` row (`app/models/llm_job.py`), status `queue`. +2. `choose_agent()` (`app/agents/agent_chooser.py`) routes the prompt. Before calling the LLM it short-circuits through `app/agents/agent_memory.py`: + - `try_handle_profile` — handles "remember X" / "who am I" locally. + - `check_prompt_cache` — exact-match cache for prior `direct_answer`/`image` responses, and a 15-minute TTL cache for `weather`. + - Only if neither hits does it call Gemini with structured output to pick an action (`car_hire_agent` / `weather_agent` / `generate_image_agent` / `marketing_agent` / `rag_ingest_agent` / `rag_query_agent` / `direct_answer` / `unsupported`). Before this call it also fetches the most recent `RagDocument` URLs (`get_known_rag_sources`) and includes them in the router's system prompt, so questions about previously-ingested sites route to `rag_query_agent` instead of `direct_answer`/`unsupported`. +3. Direct/unsupported answers finish the job immediately (status `done`). Otherwise the matching LangGraph agent is invoked with `config={"configurable": {"session": ..., "user_id": ...}}` so agent nodes can read/write `LLMMemory` (`app/models/llm_memory.py`) via the DB session. +4. If the agent graph returns `next_question` (a field on its `TypedDict` state), the job goes to status `awaiting_input` and the response is the follow-up question; the client resumes it via `PATCH /llm-job/{id}` with `{"answer": ...}`, which re-invokes the same agent (`agent_name` read back from `job.state["agent"]`) with prior state fields merged in. Otherwise the job is `done` (or `error` if the agent raised `AppException`/an exception). + +### LangGraph agents (`app/agents/`) + +Each agent (`car_hire_agent.py`, `weather_agent.py`, `generate_image_agent.py`, `marketing_agent.py`) follows the same shape: a `TypedDict` state, Pydantic schemas for structured LLM output, node functions built as `ChatPromptTemplate | llm.with_structured_output(...)`, and a compiled `StateGraph` with conditional edges for the human-in-the-loop (`missing_fields`/`next_question`) or approval (`marketing_agent`'s `approved_option`) pattern. Every LLM-calling node passes `get_langfuse_handler()` as a callback and tags calls with `langfuse_tags: [agent_name, node_name]` — follow this pattern for new nodes/agents so traces stay attributable in Langfuse. + +`rag_ingest_agent.py`/`rag_query_agent.py` are a linear (no HITL) pair implementing the URL-RAG flow from `specs/url-rag.md`: `rag_ingest_agent` regex-extracts a URL from the prompt (no LLM call — deterministic and cheaper than structured extraction for this), skips re-scraping if the URL is already in `RagDocument`, otherwise scrapes with `httpx` + `beautifulsoup4`, splits with `langchain_text_splitters.RecursiveCharacterTextSplitter`, and embeds/persists chunks via `GoogleGenerativeAIEmbeddings` (`app.config.GEMINI_EMBEDDING_MODEL`, reuses `GEMINI_API_KEY`). `rag_query_agent` embeds the question, ranks all `RagChunk` rows by cosine similarity (plain Python, no vector store — retrieval is brute-force, fine at this scale but not built to scale past a small KB), and synthesizes an answer from the top matches, falling back to a fixed "don't know" response below the similarity threshold rather than hallucinating. + +`agent_chooser.py` is the router; new agents must be registered in its system prompt (action name + routing description), re-exported from `app/agents/__init__.py`, and wired into the dispatch blocks in `app/main.py`/`routes/llm_job.py`, `cli.py`, and `mcp_server.py`. + +### Persistence + +- `app/db.py`: single SQLModel `engine` (SQLite by default), `init_db()` creates tables and seeds two demo `LLMJob` rows if empty, `get_session()` is the FastAPI dependency. +- `LLMJob` (`app/models/llm_job.py`): the job/task record — `prompt`, `response`, `status` (`queue`/`awaiting_input`/`done`/`error`), `state` (JSON blob holding the active agent name + its in-progress fields for resumption). +- `LLMMemory` (`app/models/llm_memory.py`): generic key/value cache keyed by `(user_id, memory_type, query_key)` where `query_key` is normalized via `normalize_key()` (lowercased, trimmed, trailing punctuation stripped). `memory_type` values in use: `profile`, `direct_answer`, `image`, `weather`, `weather_data`. +- `RagDocument`/`RagChunk` (`app/models/rag.py`): the URL-RAG knowledge base — one `RagDocument` per ingested URL, one `RagChunk` per text chunk with its embedding stored as JSON (`Column(JSON)`, same pattern as `LLMJob.state`). Unlike `LLMMemory`, this isn't per-`user_id` — all ingested content is shared/global across users. + +### Model/config layer + +- `app/config.py` loads all env vars (via `.env`) as module-level constants — always add new secrets/settings here rather than reading `os.environ` elsewhere. +- `app/llm_model.py` centralizes model construction (currently only `get_gemini_2_5_flash_model`, using `ChatLiteLLM`); `car_hire_agent.py` still constructs `ChatLiteLLM` directly instead of using this helper (inconsistency to be aware of, not necessarily to copy). +- `app/langfuse.py` lazily builds and caches a single `CallbackHandler`; `get_langfuse_handler()` returns `None` if keys aren't configured, and call sites are expected to handle that (`config = {...} if handler else {}`). +- `app/util.py` provides a single process-lifetime `get_session_id()` used as the Langfuse session id. + +## Repository Conventions + +From `.agents/AGENTS.md` (binding for all coding agents in this repo): + +- **Style**: PEP 8, `snake_case` functions/vars, `PascalCase` classes, `UPPER_SNAKE_CASE` constants, 4-space indents, two blank lines around top-level defs/classes, one blank line between methods. +- **Imports**: absolute imports rooted at the project (`from app.config import X`, never bare `from config import X`); expose package members via `__init__.py` `__all__` where applicable (see `app/models/__init__.py`). +- **SQLite**: never use the `COMMENT` keyword in SQL; always manage DB sessions via context managers or FastAPI `Depends`. +- **FastAPI/Pydantic**: use `lifespan` (not `@app.on_event`), Pydantic v2's `.model_dump()` (not `.dict()`), `typing.Annotated` for route dependencies. +- Read config from `app.config`, not `os.environ`, directly. +- All prompts sent to an LLM must go through Langfuse observability (pass the callback handler + `langfuse_tags`/`langfuse_session_id` metadata, matching existing node patterns). +- `TODO` comments mark code intentionally left alone — do not change unless explicitly asked. + +## Domain-specific skills + +`.agents/skills/` contains reference skills for LangChain/LangGraph/Langfuse/deep-agents patterns (fundamentals, RAG, middleware, persistence, human-in-the-loop, orchestration, etc.). Consult the relevant skill under this directory before making non-trivial changes to agent graphs, memory/persistence, or Langfuse instrumentation. diff --git a/README.md b/README.md index 6661496..0aeaa26 100644 --- a/README.md +++ b/README.md @@ -3,16 +3,28 @@ A clean, python implementation to orchestrated multi Agentic AI with LangChain/LangGraph & Langfuse (observability). ## Features -1. Multi Agents, system will chose which AI agent suitable for the tasks. +1. Multi Agents, the system will choose which AI agent is suitable for the tasks. 2. Memory capability. -3. Model agnostic, easily switch between model that suit the task. +3. Model agnostic, easily switch between models that suits the task. 4. Multi platform: CLI, MCP Server, RESTful API. +5. On-demand RAG knowledge base: add any URL and query its content afterward. ## Available Agents - `agent_chooser`: decide which agent specialized agent to use. If can answer directly (e.g: simple question, answer directly) - `agent_memory`: agent that remember the previous question and use it as base answer. Not to be confused with full RAG capability. - `weather_agent`: agent that can answer weather related question, powered by Open Weather API - `image_agent`: agent that can generate image. +- `rag_ingest_agent`: scrapes a given URL, chunks and embeds its content, and stores it in the knowledge base (SQLite). Skips re-scraping if the URL was already added. +- `rag_query_agent`: answers questions by retrieving the most relevant chunks from the knowledge base (via embedding similarity) and synthesizing an answer from them. Replies that it doesn't know if nothing relevant has been added yet. + +### RAG Knowledge Base Example +``` +You: Add https://www.kayak.co.id/ to rag knowledge base +Agent: Added 'https://www.kayak.co.id/' to the knowledge base (12 chunks indexed). + +You: What is kayak.co.id? +Agent: Kayak.co.id is a travel search engine that compares prices across... (Source: https://www.kayak.co.id/) +``` --- @@ -36,7 +48,7 @@ source venv/bin/activate touch main.db ``` -Create the `.env` file, see file [config.py](./app/config.py) for env used. +Create the `.env` file, see the file [config.py](./app/config.py) for env used. ### 3. Install Dependencies ```bash @@ -124,4 +136,6 @@ The project defines the following packages in `pyproject.toml`: 2. **`uvicorn[standard]`**: An ASGI (Asynchronous Server Gateway Interface) web server implementation for Python. FastAPI is built on ASGI standard, and `uvicorn` acts as the server to run the FastAPI application. The `[standard]` extra installs high-performance loop dependencies like `uvloop` and `httptools`. 3. **`pydantic`**: Data validation and settings management using Python type annotations. FastAPI uses Pydantic to parse and validate request JSON payloads, and serialize response objects. 4. **`pytest`**: A robust testing framework for writing clean, readable, and scalable unit tests. -5. **`httpx`**: A next-generation HTTP client for Python. It is used in unit tests alongside FastAPI's `TestClient` to make asynchronous/synchronous HTTP requests to the application. \ No newline at end of file +5. **`httpx`**: A next-generation HTTP client for Python. It is used in unit tests alongside FastAPI's `TestClient` to make asynchronous/synchronous HTTP requests to the application. Also used by `rag_ingest_agent` to fetch page content. +6. **`beautifulsoup4`**: HTML parser used by `rag_ingest_agent` to strip tags and extract readable text from scraped pages. +7. **`langchain-text-splitters`**: Splits scraped page text into overlapping chunks before embedding, for the RAG knowledge base. \ No newline at end of file diff --git a/app/agents/__init__.py b/app/agents/__init__.py index dcba0cc..45a1511 100644 --- a/app/agents/__init__.py +++ b/app/agents/__init__.py @@ -3,3 +3,5 @@ from .agent_chooser import choose_agent from .generate_image_agent import generate_image_agent from .marketing_agent import marketing_agent +from .rag_ingest_agent import rag_ingest_agent +from .rag_query_agent import rag_query_agent diff --git a/app/agents/agent_chooser.py b/app/agents/agent_chooser.py index fb9cefd..d16dc53 100644 --- a/app/agents/agent_chooser.py +++ b/app/agents/agent_chooser.py @@ -1,21 +1,22 @@ from app.llm_model import get_gemini_2_5_flash_model from app.config import GEMINI_API_KEY from pydantic import BaseModel, Field -from typing import Optional +from typing import List, Optional import os from langchain_core.prompts import ChatPromptTemplate from langchain_litellm import ChatLiteLLM from app.log import get_logger from app.langfuse import get_langfuse_handler from app.util import get_session_id -from sqlmodel import Session +from sqlmodel import Session, select, col from app.agents.agent_memory import try_handle_profile, check_prompt_cache, save_memory +from app.models import RagDocument class ChooserResult(BaseModel): action: str = Field( ..., - description="The action to take: 'car_hire_agent', 'weather_agent', 'generate_image_agent', 'marketing_agent', 'direct_answer', or 'unsupported'" + description="The action to take: 'car_hire_agent', 'weather_agent', 'generate_image_agent', 'marketing_agent', 'rag_ingest_agent', 'rag_query_agent', 'direct_answer', or 'unsupported'" ) direct_response: Optional[str] = Field( None, @@ -25,6 +26,18 @@ class ChooserResult(BaseModel): log = get_logger('agent-chooser') +KNOWN_SOURCES_LIMIT = 20 + + +def get_known_rag_sources(session: Session, limit: int = KNOWN_SOURCES_LIMIT) -> List[str]: + """ + Fetch a bounded list of previously ingested URLs so the router can recognize + questions that relate to content already added to the RAG knowledge base. + """ + statement = select(RagDocument).order_by(col(RagDocument.fetched_at).desc()).limit(limit) + documents = session.exec(statement).all() + return [doc.url for doc in documents] + def choose_agent(prompt: str, session: Optional[Session] = None, user_id: str = "default_user") -> ChooserResult: """ @@ -58,6 +71,12 @@ def choose_agent(prompt: str, session: Optional[Session] = None, user_id: str = structured_llm = llm.with_structured_output(ChooserResult) + known_sources = get_known_rag_sources(session) if session else [] + known_sources_text = ( + "\n".join(f"- {url}" for url in known_sources) + if known_sources else "(none yet)" + ) + # TODO use abstraction to build system prompt and agent discovery from app.agents. prompt_template = ChatPromptTemplate.from_messages([ ("system", ( @@ -66,9 +85,13 @@ def choose_agent(prompt: str, session: Optional[Session] = None, user_id: str = "- If the user asks about the weather, temperature, rain, or weather forecast, set action to 'weather_agent'.\n" "- If the user asks to generate, draw, paint, create, or design an image, picture, photo, or graphic, set action to 'generate_image_agent'.\n" "- If the user asks to draft, generate, or write a marketing campaign, ad copy, marketing copy, or Facebook/Instagram/Google ad, set action to 'marketing_agent'.\n" + "- If the user asks to add, scrape, index, or learn a URL/website into the knowledge base (e.g. 'add https://... to rag knowledge base'), set action to 'rag_ingest_agent'.\n" + "- If the user asks a question that relates to one of the Known Knowledge Base Sources below (by domain, brand, or topic, even if not an exact match), or explicitly asks to look something up 'from the knowledge base' / 'from what you scraped', set action to 'rag_query_agent'.\n" "- If the user asks to remember a fact/preference about them, or asks what you know/remember about them (e.g. 'who am i', 'do you know my name'), set action to 'direct_answer' and write a helpful response in direct_response.\n" "- If the user asks a general knowledge question that you can answer directly (e.g. 'what is E=mc2'), set action to 'direct_answer' and write the answer in direct_response.\n" - "- If the user asks for anything else that requires a tool we do not have (e.g. booking a flight, ordering food, writing code), set action to 'unsupported' and set direct_response to exactly 'I can not do that at the moment'." + "- If the user asks for anything else that requires a tool we do not have (e.g. booking a flight, ordering food, writing code), set action to 'unsupported' and set direct_response to exactly 'I can not do that at the moment'.\n\n" + "Known Knowledge Base Sources (URLs already added via rag_ingest_agent):\n" + "{known_sources}" )), ("user", "{prompt}") ]) @@ -84,7 +107,7 @@ def choose_agent(prompt: str, session: Optional[Session] = None, user_id: str = } } if handler else {} - result = chain.invoke({"prompt": prompt}, config=config) + result = chain.invoke({"prompt": prompt, "known_sources": known_sources_text}, config=config) if result.action == "unsupported": result.direct_response = "I can not do that at the moment" diff --git a/app/agents/rag_ingest_agent.py b/app/agents/rag_ingest_agent.py new file mode 100644 index 0000000..f60ba1c --- /dev/null +++ b/app/agents/rag_ingest_agent.py @@ -0,0 +1,182 @@ +import re +from typing import TypedDict, Optional, Dict, Any +import httpx +from bs4 import BeautifulSoup +from sqlmodel import select +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_google_genai import GoogleGenerativeAIEmbeddings +from langchain_core.runnables import RunnableConfig +from langgraph.graph import StateGraph, END +from app.app_exception import AppException +from app.config import GEMINI_API_KEY, GEMINI_EMBEDDING_MODEL +from app.models import RagDocument, RagChunk +from app.log import get_logger + +log = get_logger("rag-ingest-agent") + +URL_REGEX = re.compile(r"https?://[^\s<>\"')\]]+") +REQUEST_TIMEOUT = 10.0 +MAX_CONTENT_CHARS = 300_000 +CHUNK_SIZE = 1000 +CHUNK_OVERLAP = 200 + + +class RagIngestState(TypedDict): + prompt: str + url: Optional[str] + title: Optional[str] + scraped_text: Optional[str] + already_exists: Optional[bool] + chunk_count: Optional[int] + final_response: Optional[str] + + +def _get_session(config: Optional[RunnableConfig]): + if config and isinstance(config, dict) and config.get("configurable"): + return config["configurable"].get("session") + return None + + +# Node 1: Extract the URL from the prompt (plain regex, no LLM call needed) +def extract_url(state: RagIngestState) -> Dict[str, Any]: + match = URL_REGEX.search(state["prompt"]) + if not match: + return { + "url": None, + "final_response": "Please provide a valid http(s) URL to add to the knowledge base." + } + url = match.group(0).rstrip(".,;:!?") + if not url.lower().startswith(("http://", "https://")): + return { + "url": None, + "final_response": "Please provide a valid http(s) URL to add to the knowledge base." + } + return {"url": url} + + +# Node 2: Skip re-scraping if the URL is already indexed +def check_existing(state: RagIngestState, config: RunnableConfig = None) -> Dict[str, Any]: + session = _get_session(config) + if session: + existing = session.exec(select(RagDocument).where(RagDocument.url == state["url"])).first() + if existing: + log.info(f"URL already indexed, skipping scrape: {state['url']}") + return { + "already_exists": True, + "final_response": f"'{state['url']}' is already in the knowledge base ({existing.char_count} characters indexed)." + } + return {"already_exists": False} + + +# Node 3: Scrape the URL and extract readable text +def scrape_url(state: RagIngestState) -> Dict[str, Any]: + url = state["url"] + log.info(f"Scraping {url}") + try: + headers = {"User-Agent": "Mozilla/5.0 (compatible; AgentOrchestratorBot/1.0)"} + response = httpx.get(url, timeout=REQUEST_TIMEOUT, headers=headers, follow_redirects=True) + response.raise_for_status() + except Exception as e: + raise AppException(f"failed to fetch {url}: {str(e)}") + + soup = BeautifulSoup(response.text, "html.parser") + for tag in soup(["script", "style", "noscript"]): + tag.decompose() + + title = soup.title.string.strip() if soup.title and soup.title.string else url + lines = [line.strip() for line in soup.get_text(separator="\n").splitlines() if line.strip()] + cleaned = "\n".join(lines)[:MAX_CONTENT_CHARS] + + if not cleaned: + raise AppException(f"no readable content found at {url}") + + return {"scraped_text": cleaned, "title": title} + + +# Node 4: Split into chunks, embed, and persist +def chunk_and_embed(state: RagIngestState, config: RunnableConfig = None) -> Dict[str, Any]: + if not GEMINI_API_KEY: + raise AppException("please set GEMINI_API_KEY in env") + + session = _get_session(config) + if not session: + raise AppException("no database session available to store knowledge base content") + + splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP) + chunks = splitter.split_text(state["scraped_text"]) + + if not chunks: + raise AppException(f"no content to index from {state['url']}") + + embeddings_model = GoogleGenerativeAIEmbeddings(model=GEMINI_EMBEDDING_MODEL, google_api_key=GEMINI_API_KEY) + try: + vectors = embeddings_model.embed_documents(chunks, task_type="RETRIEVAL_DOCUMENT") + except Exception as e: + raise AppException(f"failed to embed content: {str(e)}") + + document = RagDocument(url=state["url"], title=state.get("title"), char_count=len(state["scraped_text"])) + session.add(document) + session.commit() + session.refresh(document) + + for idx, (chunk_text, vector) in enumerate(zip(chunks, vectors)): + session.add(RagChunk(document_id=document.id, chunk_index=idx, content=chunk_text, embedding=vector)) + session.commit() + + log.info(f"Indexed {len(chunks)} chunks for {state['url']}") + return {"chunk_count": len(chunks)} + + +# Node 5: Format the final response +def format_response(state: RagIngestState) -> Dict[str, Any]: + response = ( + f"Added '{state['url']}' to the knowledge base " + f"({state.get('chunk_count', 0)} chunks indexed)." + ) + return {"final_response": response} + + +# Conditional routing +def route_after_extract(state: RagIngestState) -> str: + return "check_existing" if state.get("url") else "no_url" + + +def route_after_check(state: RagIngestState) -> str: + return "exists" if state.get("already_exists") else "scrape" + + +# Construct the graph +workflow = StateGraph(RagIngestState) + +workflow.add_node("extract_url", extract_url) +workflow.add_node("check_existing", check_existing) +workflow.add_node("scrape_url", scrape_url) +workflow.add_node("chunk_and_embed", chunk_and_embed) +workflow.add_node("format_response", format_response) + +workflow.set_entry_point("extract_url") + +workflow.add_conditional_edges( + "extract_url", + route_after_extract, + { + "no_url": END, + "check_existing": "check_existing" + } +) + +workflow.add_conditional_edges( + "check_existing", + route_after_check, + { + "exists": END, + "scrape": "scrape_url" + } +) + +workflow.add_edge("scrape_url", "chunk_and_embed") +workflow.add_edge("chunk_and_embed", "format_response") +workflow.add_edge("format_response", END) + +# Compile graph +rag_ingest_agent = workflow.compile() diff --git a/app/agents/rag_query_agent.py b/app/agents/rag_query_agent.py new file mode 100644 index 0000000..a376419 --- /dev/null +++ b/app/agents/rag_query_agent.py @@ -0,0 +1,138 @@ +from typing import TypedDict, Optional, List, Dict, Any +from sqlmodel import select +from langchain_core.prompts import ChatPromptTemplate +from langchain_core.runnables import RunnableConfig +from langchain_google_genai import GoogleGenerativeAIEmbeddings +from langgraph.graph import StateGraph, END +from app.app_exception import AppException +from app.llm_model import get_gemini_2_5_flash_model +from app.config import GEMINI_API_KEY, GEMINI_EMBEDDING_MODEL +from app.models import RagChunk, RagDocument +from app.log import get_logger +from app.langfuse import get_langfuse_handler +from app.util import get_session_id + +log = get_logger("rag-query-agent") + +TOP_K = 5 +MIN_SIMILARITY = 0.6 +NO_KNOWLEDGE_RESPONSE = "I don't have information about that in my knowledge base yet. Try adding a relevant URL first." + + +class RagQueryState(TypedDict): + prompt: str + query_embedding: Optional[List[float]] + retrieved_chunks: Optional[List[Dict[str, Any]]] + final_response: Optional[str] + + +def _get_session(config: Optional[RunnableConfig]): + if config and isinstance(config, dict) and config.get("configurable"): + return config["configurable"].get("session") + return None + + +def _cosine_similarity(a: List[float], b: List[float]) -> float: + dot = sum(x * y for x, y in zip(a, b)) + norm_a = sum(x * x for x in a) ** 0.5 + norm_b = sum(y * y for y in b) ** 0.5 + if norm_a == 0 or norm_b == 0: + return 0.0 + return dot / (norm_a * norm_b) + + +# Node 1: Embed the user's question +def embed_query(state: RagQueryState) -> Dict[str, Any]: + if not GEMINI_API_KEY: + raise AppException("please set GEMINI_API_KEY in env") + + embeddings_model = GoogleGenerativeAIEmbeddings(model=GEMINI_EMBEDDING_MODEL, google_api_key=GEMINI_API_KEY) + try: + vector = embeddings_model.embed_query(state["prompt"], task_type="RETRIEVAL_QUERY") + except Exception as e: + raise AppException(f"failed to embed query: {str(e)}") + + return {"query_embedding": vector} + + +# Node 2: Rank stored chunks by cosine similarity to the query +def retrieve_chunks(state: RagQueryState, config: RunnableConfig = None) -> Dict[str, Any]: + session = _get_session(config) + if not session: + return {"retrieved_chunks": []} + + query_vector = state["query_embedding"] + all_chunks = session.exec(select(RagChunk)).all() + + scored = [] + for chunk in all_chunks: + if not chunk.embedding: + continue + score = _cosine_similarity(query_vector, chunk.embedding) + if score >= MIN_SIMILARITY: + scored.append((score, chunk)) + + scored.sort(key=lambda pair: pair[0], reverse=True) + + results = [] + for score, chunk in scored[:TOP_K]: + document = session.get(RagDocument, chunk.document_id) + results.append({ + "content": chunk.content, + "url": document.url if document else None, + "score": score + }) + + log.info(f"Retrieved {len(results)} relevant chunks for query") + return {"retrieved_chunks": results} + + +# Node 3: Synthesize the answer from retrieved context +def synthesize_answer(state: RagQueryState) -> Dict[str, Any]: + retrieved = state.get("retrieved_chunks") or [] + if not retrieved: + return {"final_response": NO_KNOWLEDGE_RESPONSE} + + context = "\n\n".join(f"Source: {chunk['url']}\n{chunk['content']}" for chunk in retrieved) + + llm = get_gemini_2_5_flash_model(temperature=0.2) + prompt_template = ChatPromptTemplate.from_messages([ + ("system", ( + "You are a helpful assistant answering questions using ONLY the provided knowledge base context.\n" + "If the context does not contain the answer, say you don't have that information.\n" + "Cite the source URL(s) you used." + )), + ("user", "Context:\n{context}\n\nQuestion: {prompt}") + ]) + chain = prompt_template | llm + + try: + handler = get_langfuse_handler() + config = { + "callbacks": [handler], + "metadata": { + "langfuse_session_id": get_session_id(), + "langfuse_tags": ["rag_query_agent", "synthesize_answer"] + } + } if handler else {} + response = chain.invoke({"context": context, "prompt": state["prompt"]}, config=config).content + except Exception as e: + raise AppException(f"failed to synthesize answer: {str(e)}") + + return {"final_response": response} + + +# Construct the graph +workflow = StateGraph(RagQueryState) + +workflow.add_node("embed_query", embed_query) +workflow.add_node("retrieve_chunks", retrieve_chunks) +workflow.add_node("synthesize_answer", synthesize_answer) + +workflow.set_entry_point("embed_query") +workflow.add_edge("embed_query", "retrieve_chunks") +workflow.add_edge("retrieve_chunks", "synthesize_answer") +workflow.add_edge("synthesize_answer", END) + +# Compile graph +rag_query_agent = workflow.compile() diff --git a/app/config.py b/app/config.py index 95b57f0..b44c8f2 100644 --- a/app/config.py +++ b/app/config.py @@ -7,6 +7,7 @@ DATABASE_URL = os.getenv("DATABASE_URL") GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") +GEMINI_EMBEDDING_MODEL = os.getenv("GEMINI_EMBEDDING_MODEL", "gemini-embedding-2-preview") LANGFUSE_SECRET_KEY = os.getenv("LANGFUSE_SECRET_KEY") LANGFUSE_PUBLIC_KEY = os.getenv("LANGFUSE_PUBLIC_KEY") LANGFUSE_BASE_URL = os.getenv("LANGFUSE_BASE_URL") diff --git a/app/models/__init__.py b/app/models/__init__.py index 9cdd404..90c2fca 100644 --- a/app/models/__init__.py +++ b/app/models/__init__.py @@ -1,5 +1,6 @@ from app.models.llm_job import LLMJob from app.models.llm_memory import LLMMemory +from app.models.rag import RagDocument, RagChunk -__all__ = ["LLMJob", "LLMMemory"] +__all__ = ["LLMJob", "LLMMemory", "RagDocument", "RagChunk"] diff --git a/app/models/rag.py b/app/models/rag.py new file mode 100644 index 0000000..5a1f43d --- /dev/null +++ b/app/models/rag.py @@ -0,0 +1,23 @@ +from datetime import datetime +from sqlmodel import Field, SQLModel, Column, JSON + + +class RagDocument(SQLModel, table=True): + __tablename__ = "rag_documents" + + id: int | None = Field(default=None, primary_key=True) + url: str = Field(nullable=False, index=True) + title: str | None = Field(default=None) + char_count: int = Field(default=0) + fetched_at: datetime = Field(default_factory=datetime.utcnow) + + +class RagChunk(SQLModel, table=True): + __tablename__ = "rag_chunks" + + id: int | None = Field(default=None, primary_key=True) + document_id: int = Field(foreign_key="rag_documents.id", nullable=False, index=True) + chunk_index: int = Field(nullable=False) + content: str = Field(nullable=False) + embedding: list[float] = Field(default=None, sa_column=Column(JSON)) + created_at: datetime = Field(default_factory=datetime.utcnow) diff --git a/app/routes/llm_job.py b/app/routes/llm_job.py index a8c7350..5d2632d 100644 --- a/app/routes/llm_job.py +++ b/app/routes/llm_job.py @@ -5,7 +5,7 @@ from pydantic import BaseModel, Field from app.db import get_session from app.models import LLMJob -from app.agents import car_hire_agent, weather_agent, choose_agent, generate_image_agent +from app.agents import car_hire_agent, weather_agent, choose_agent, generate_image_agent, rag_ingest_agent, rag_query_agent from app.langfuse import get_langfuse_handler @@ -108,6 +108,49 @@ def create_job(payload: LLMJobCreate, session: Session = Depends(get_session)): job.response = f"Image generation agent failed: {str(e)}" job.responded_at = datetime.utcnow() job.state = {"agent": "generate_image_agent"} + elif choice.action == "rag_ingest_agent": + initial_state = { + "prompt": payload.prompt, + "url": None, + "title": None, + "scraped_text": None, + "already_exists": None, + "chunk_count": None, + "final_response": None + } + try: + result = rag_ingest_agent.invoke(initial_state, config=config) + job.status = "done" + job.response = result.get("final_response") or "Done" + job.responded_at = datetime.utcnow() + job.state = { + "agent": "rag_ingest_agent", + "url": result.get("url"), + "chunk_count": result.get("chunk_count") + } + except Exception as e: + job.status = "error" + job.response = f"RAG ingest agent failed: {str(e)}" + job.responded_at = datetime.utcnow() + job.state = {"agent": "rag_ingest_agent"} + elif choice.action == "rag_query_agent": + initial_state = { + "prompt": payload.prompt, + "query_embedding": None, + "retrieved_chunks": None, + "final_response": None + } + try: + result = rag_query_agent.invoke(initial_state, config=config) + job.status = "done" + job.response = result.get("final_response") or "Done" + job.responded_at = datetime.utcnow() + job.state = {"agent": "rag_query_agent"} + except Exception as e: + job.status = "error" + job.response = f"RAG query agent failed: {str(e)}" + job.responded_at = datetime.utcnow() + job.state = {"agent": "rag_query_agent"} else: # Default to car_hire_agent diff --git a/cli.py b/cli.py index 1a4138f..7c24eea 100644 --- a/cli.py +++ b/cli.py @@ -1,7 +1,7 @@ from app.app_exception import AppException from app.log import get_logger import sys -from app.agents import choose_agent, car_hire_agent, weather_agent, generate_image_agent, marketing_agent +from app.agents import choose_agent, car_hire_agent, weather_agent, generate_image_agent, marketing_agent, rag_ingest_agent, rag_query_agent from sqlmodel import Session from app.db import engine, init_db @@ -154,6 +154,45 @@ def run_marketing_flow(initial_prompt: str, session: Session = None) -> None: print(f"\nAgent: {result.get('final_response') or 'Campaign finalized.'}") +def run_rag_ingest_flow(initial_prompt: str, session: Session = None) -> None: + state = { + "prompt": initial_prompt, + "url": None, + "title": None, + "scraped_text": None, + "already_exists": None, + "chunk_count": None, + "final_response": None + } + config = { + "configurable": { + "session": session, + "user_id": "default_user" + } + } if session else {} + + result = rag_ingest_agent.invoke(state, config=config) + print(f"\nAgent: {result.get('final_response') or 'Done.'}") + + +def run_rag_query_flow(initial_prompt: str, session: Session = None) -> None: + state = { + "prompt": initial_prompt, + "query_embedding": None, + "retrieved_chunks": None, + "final_response": None + } + config = { + "configurable": { + "session": session, + "user_id": "default_user" + } + } if session else {} + + result = rag_query_agent.invoke(state, config=config) + print(f"\nAgent: {result.get('final_response') or 'Done.'}") + + def main() -> None: init_db() print("==================================================") @@ -194,6 +233,12 @@ def main() -> None: elif choice.action == "marketing_agent": print("\n[Routing to Marketing Agent...]") run_marketing_flow(prompt, session=session) + elif choice.action == "rag_ingest_agent": + print("\n[Routing to RAG Ingest Agent...]") + run_rag_ingest_flow(prompt, session=session) + elif choice.action == "rag_query_agent": + print("\n[Routing to RAG Query Agent...]") + run_rag_query_flow(prompt, session=session) elif choice.action == "unsupported": print(f"\nAgent: {choice.direct_response}") else: diff --git a/mcp_server.py b/mcp_server.py index 66a2ddf..7212f43 100644 --- a/mcp_server.py +++ b/mcp_server.py @@ -4,7 +4,7 @@ from sqlmodel import Session from app.db import engine, init_db from app.models import LLMJob -from app.agents import choose_agent, weather_agent, generate_image_agent, marketing_agent +from app.agents import choose_agent, weather_agent, generate_image_agent, marketing_agent, rag_ingest_agent, rag_query_agent # 1. Initialize the FastMCP Server @@ -79,6 +79,25 @@ def submit_agent_task(prompt: str) -> str: "next_question": None } result = marketing_agent.invoke(initial_state, config=config) + elif choice.action == "rag_ingest_agent": + initial_state = { + "prompt": prompt, + "url": None, + "title": None, + "scraped_text": None, + "already_exists": None, + "chunk_count": None, + "final_response": None + } + result = rag_ingest_agent.invoke(initial_state, config=config) + elif choice.action == "rag_query_agent": + initial_state = { + "prompt": prompt, + "query_embedding": None, + "retrieved_chunks": None, + "final_response": None + } + result = rag_query_agent.invoke(initial_state, config=config) else: # Unsupported action inside this MCP context job.status = "error" @@ -100,7 +119,7 @@ def submit_agent_task(prompt: str) -> str: job.state = { "agent": choice.action, - **{k: v for k, v in result.items() if k not in ("prompt", "final_response", "next_question") and not isinstance(v, bytes)} + **{k: v for k, v in result.items() if k not in ("prompt", "final_response", "next_question", "query_embedding", "scraped_text") and not isinstance(v, bytes)} } session.add(job) session.commit() diff --git a/pyproject.toml b/pyproject.toml index 578bed1..7afb570 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,4 +22,6 @@ dependencies = [ "langchain>=1.3.12", "boto3>=1.34.0", "mcp>=0.1.0", + "beautifulsoup4>=4.12.0", + "langchain-text-splitters>=0.3.0", ] diff --git a/specs/url-rag.md b/specs/url-rag.md new file mode 100644 index 0000000..a0699f5 --- /dev/null +++ b/specs/url-rag.md @@ -0,0 +1,14 @@ +Goals: +Have an RAG agent that can add knowledge based on URL given and then can be queried by prompt. + +For example, prompt: "Add https://www.kayak.co.id/ to rag knowledge base" + +This will make agent_chooser to call the RAG agent to add the URL to the knowledge base. +The RAG agent will then scrape the website and add the content to the knowledge base. +The information is saved in the sqlite database and can be queried by the RAG agent. + +Then you can query the knowledge base by prompt: +"What is kayak.co.id?" and the RAG agent will return the answer based on the knowledge it has added from the URL. + +Another prompt: +"How kayak find such a low car rental price?" and the RAG agent will return the answer based on the knowledge it has added from the URL. \ No newline at end of file diff --git a/tests/test_rag_agents.py b/tests/test_rag_agents.py new file mode 100644 index 0000000..0d51731 --- /dev/null +++ b/tests/test_rag_agents.py @@ -0,0 +1,176 @@ +import pytest +from unittest.mock import patch, MagicMock +from sqlmodel import SQLModel, create_engine, Session, StaticPool, select + +from app.models import RagDocument, RagChunk +from app.agents.rag_ingest_agent import rag_ingest_agent +from app.agents.rag_query_agent import rag_query_agent, NO_KNOWLEDGE_RESPONSE + +# Create an in-memory SQLite engine +engine = create_engine( + "sqlite://", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, +) + + +@pytest.fixture(autouse=True) +def setup_and_teardown_db(): + SQLModel.metadata.create_all(engine) + yield + SQLModel.metadata.drop_all(engine) + + +MOCK_HTML = """ + +
This domain is for use in illustrative examples.
+ + +""" + + +def fake_embed_documents(texts, task_type=None): + return [[1.0, 0.0] for _ in texts] + + +def test_rag_ingest_agent_success(): + mock_http_response = MagicMock() + mock_http_response.text = MOCK_HTML + mock_http_response.raise_for_status.return_value = None + + with Session(engine) as session, \ + patch("httpx.get", return_value=mock_http_response), \ + patch("app.agents.rag_ingest_agent.GoogleGenerativeAIEmbeddings") as mock_embeddings_cls: + + mock_embeddings_instance = MagicMock() + mock_embeddings_instance.embed_documents.side_effect = fake_embed_documents + mock_embeddings_cls.return_value = mock_embeddings_instance + + state = { + "prompt": "Add https://www.example.com/ to rag knowledge base", + "url": None, + "title": None, + "scraped_text": None, + "already_exists": None, + "chunk_count": None, + "final_response": None + } + config = {"configurable": {"session": session, "user_id": "default_user"}} + + result = rag_ingest_agent.invoke(state, config=config) + + assert result["url"] == "https://www.example.com/" + assert result["chunk_count"] == 1 + assert "Added" in result["final_response"] + + documents = session.exec(select(RagDocument)).all() + chunks = session.exec(select(RagChunk)).all() + assert len(documents) == 1 + assert documents[0].url == "https://www.example.com/" + assert len(chunks) == 1 + assert chunks[0].embedding == [1.0, 0.0] + + +def test_rag_ingest_agent_already_exists(): + with Session(engine) as session: + session.add(RagDocument(url="https://www.example.com/", title="Example", char_count=100)) + session.commit() + + with Session(engine) as session, patch("httpx.get") as mock_get: + state = { + "prompt": "Add https://www.example.com/ to rag knowledge base", + "url": None, + "title": None, + "scraped_text": None, + "already_exists": None, + "chunk_count": None, + "final_response": None + } + config = {"configurable": {"session": session, "user_id": "default_user"}} + + result = rag_ingest_agent.invoke(state, config=config) + + assert "already in the knowledge base" in result["final_response"] + mock_get.assert_not_called() + + +def test_rag_ingest_agent_no_url(): + with Session(engine) as session: + state = { + "prompt": "Add kayak.co.id to rag knowledge base", + "url": None, + "title": None, + "scraped_text": None, + "already_exists": None, + "chunk_count": None, + "final_response": None + } + config = {"configurable": {"session": session, "user_id": "default_user"}} + + result = rag_ingest_agent.invoke(state, config=config) + + assert "valid http(s) URL" in result["final_response"] + + +def test_rag_query_agent_no_knowledge(): + with Session(engine) as session, \ + patch("app.agents.rag_query_agent.GoogleGenerativeAIEmbeddings") as mock_embeddings_cls: + + mock_embeddings_instance = MagicMock() + mock_embeddings_instance.embed_query.return_value = [1.0, 0.0] + mock_embeddings_cls.return_value = mock_embeddings_instance + + state = { + "prompt": "What is kayak.co.id?", + "query_embedding": None, + "retrieved_chunks": None, + "final_response": None + } + config = {"configurable": {"session": session, "user_id": "default_user"}} + + result = rag_query_agent.invoke(state, config=config) + + assert result["final_response"] == NO_KNOWLEDGE_RESPONSE + + +def test_rag_query_agent_with_match(): + with Session(engine) as session: + doc = RagDocument(url="https://www.kayak.co.id/", title="Kayak", char_count=500) + session.add(doc) + session.commit() + session.refresh(doc) + session.add(RagChunk(document_id=doc.id, chunk_index=0, content="Kayak compares flight and car prices.", embedding=[1.0, 0.0])) + session.commit() + + with Session(engine) as session, \ + patch("app.agents.rag_query_agent.GoogleGenerativeAIEmbeddings") as mock_embeddings_cls, \ + patch("app.agents.rag_query_agent.get_gemini_2_5_flash_model") as mock_llm_factory: + + mock_embeddings_instance = MagicMock() + mock_embeddings_instance.embed_query.return_value = [1.0, 0.0] + mock_embeddings_cls.return_value = mock_embeddings_instance + + mock_response = MagicMock() + mock_response.content = "Kayak compares prices across providers. Source: https://www.kayak.co.id/" + mock_llm = MagicMock() + mock_llm.invoke.return_value = mock_response + mock_llm.return_value = mock_response + mock_llm_factory.return_value = mock_llm + + state = { + "prompt": "What is kayak.co.id?", + "query_embedding": None, + "retrieved_chunks": None, + "final_response": None + } + config = {"configurable": {"session": session, "user_id": "default_user"}} + + result = rag_query_agent.invoke(state, config=config) + + assert len(result["retrieved_chunks"]) == 1 + assert result["retrieved_chunks"][0]["url"] == "https://www.kayak.co.id/" + assert result["final_response"] == mock_response.content diff --git a/uv.lock b/uv.lock index f8084c9..5f23ff4 100644 --- a/uv.lock +++ b/uv.lock @@ -199,6 +199,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/df/73/b6e24bd22e6720ca8ee9a85a0c4a2971af8497d8f3193fa05390cbd46e09/backoff-2.2.1-py3-none-any.whl", hash = "sha256:63579f9a0628e06278f7e47b7d7d5b6ce20dc65c5e96a6f3ca99a6adca0396e8", size = 15148, upload-time = "2022-10-05T19:19:30.546Z" }, ] +[[package]] +name = "beautifulsoup4" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "soupsieve" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/65/318323f98dbee45d42dff61d8f047181bc6f2268a9068cfad035a46be5af/beautifulsoup4-4.15.0.tar.gz", hash = "sha256:288e3ca7d54b06f2ac191970bc275c1939cb46d450b255bf6718b04aa37ab4f7", size = 632571, upload-time = "2026-06-07T16:44:20.453Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/c6/92fcd42f1ba33e1184263f25bfabf3d27c383410470f169e4b8163bf9c17/beautifulsoup4-4.15.0-py3-none-any.whl", hash = "sha256:d6f88de62e1d4e38ecb1077eb9724cd0eff29d2a08ca16a401e9b9e93f117cf9", size = 109924, upload-time = "2026-06-07T16:44:21.566Z" }, +] + [[package]] name = "boto3" version = "1.43.46" @@ -2334,6 +2347,7 @@ name = "python-llm-api" version = "0.1.0" source = { virtual = "." } dependencies = [ + { name = "beautifulsoup4" }, { name = "boto3" }, { name = "fastapi" }, { name = "httpx" }, @@ -2342,6 +2356,7 @@ dependencies = [ { name = "langchain-core" }, { name = "langchain-google-genai" }, { name = "langchain-litellm" }, + { name = "langchain-text-splitters" }, { name = "langfuse" }, { name = "langgraph" }, { name = "litellm" }, @@ -2355,6 +2370,7 @@ dependencies = [ [package.metadata] requires-dist = [ + { name = "beautifulsoup4", specifier = ">=4.12.0" }, { name = "boto3", specifier = ">=1.34.0" }, { name = "fastapi", specifier = ">=0.111.0" }, { name = "httpx", specifier = ">=0.27.0" }, @@ -2363,6 +2379,7 @@ requires-dist = [ { name = "langchain-core", specifier = ">=0.2.0" }, { name = "langchain-google-genai", specifier = ">=1.0.0" }, { name = "langchain-litellm", specifier = ">=0.6.4" }, + { name = "langchain-text-splitters", specifier = ">=0.3.0" }, { name = "langfuse", specifier = ">=2.0.0" }, { name = "langgraph", specifier = ">=0.1.0" }, { name = "litellm", specifier = ">=1.83.7" }, @@ -2758,6 +2775,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" }, ] +[[package]] +name = "soupsieve" +version = "2.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/80/f1/93422647dd7e461f23d254e6b2bfa687a85b53aeb4903fcdbb74474d4584/soupsieve-2.9.tar.gz", hash = "sha256:acee8417325c5653e1377dc31eccad59eb82cbc65942afe6174c53b3aaad63fc", size = 122122, upload-time = "2026-07-19T01:35:18.425Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/d6/3185ab5ad1280319b31986898f3206dd7227cd75e293d4dba2a5e6bf27a0/soupsieve-2.9-py3-none-any.whl", hash = "sha256:a2b2c76d67df2382d245409fd71e321a571717e58463efa32ace87dcadac2c12", size = 37387, upload-time = "2026-07-19T01:35:17.106Z" }, +] + [[package]] name = "sqlalchemy" version = "2.0.51"