diff --git a/uc-0a/agents.md b/uc-0a/agents.md index cd4d882..a6eefae 100644 --- a/uc-0a/agents.md +++ b/uc-0a/agents.md @@ -1,27 +1,15 @@ -# agents.md — UC-0A Complaint Classifier -# INSTRUCTIONS: -# 1. Open your AI tool -# 2. Paste the full contents of uc-0a/README.md -# 3. Use this prompt: -# "Read this UC README. Using the R.I.C.E framework, generate an -# agents.md YAML with four fields: role, intent, context, enforcement. -# Enforcement must include every rule listed under -# 'Enforcement Rules Your agents.md Must Include'. -# Output only valid YAML." -# 4. Paste the output below - role: > - [FILL IN] + An AI classifier for the City Operations team that reads municipal complaints and categorizes them. Its output directly feeds the Director's dashboard. intent: > - [FILL IN] + To accurately classify each complaint into a predefined category, assign a priority level, provide a specific reason for the classification, and flag ambiguous complaints for review. context: > - [FILL IN] + The agent must only use the provided complaint row (specifically the description and location fields) to make its determination. It must not use external knowledge or invent details not present in the complaint. enforcement: - - "[FILL IN: category enum rule]" - - "[FILL IN: severity keyword rule — list the keywords]" - - "[FILL IN: reason field rule]" - - "[FILL IN: ambiguity refusal rule]" - - "[FILL IN: no invented categories rule]" + - "Category must be exactly one value from the allowed list: Pothole, Flooding, Streetlight, Waste, Noise, Road Damage, Heritage Damage, Heat Hazard, Drain Blockage, Other. No variations." + - "Priority must be Urgent if description contains any severity keyword: injury, child, school, hospital, ambulance, fire, hazard, fell, collapse." + - "Every output row must include a reason field citing specific words from the description." + - "If category cannot be determined confidently — output `category: Other` and `flag: NEEDS_REVIEW`." + - "Never invent category names outside the allowed list." diff --git a/uc-0a/classifier.py b/uc-0a/classifier.py index 3f8fe55..8659b0e 100644 --- a/uc-0a/classifier.py +++ b/uc-0a/classifier.py @@ -1,26 +1,137 @@ -""" -UC-0A — Complaint Classifier -classifier.py — Starter file - -Build this using your AI coding tool: -1. Share agents.md, skills.md, and uc-0a/README.md -2. Ask the AI to implement this file -3. Run: python3 classifier.py --input ../data/city-test-files/test_pune.csv \ - --output results_pune.csv -""" import argparse import csv +import os +import json +import time + +def call_llm(prompt: str) -> str: + """Uses Gemini API to classify the complaint based on the prompt.""" + api_key = os.environ.get("GEMINI_API_KEY") + if not api_key: + return json.dumps({ + "category": "Other", + "priority": "Standard", + "reason": "[ERROR] GEMINI_API_KEY not set", + "flag": "NEEDS_REVIEW" + }) + + try: + import google.generativeai as genai + genai.configure(api_key=api_key) + # We request JSON response from Gemini + model = genai.GenerativeModel( + "gemini-1.5-flash", + generation_config={"response_mime_type": "application/json"} + ) + response = model.generate_content(prompt) + return response.text + except Exception as e: + return json.dumps({ + "category": "Other", + "priority": "Standard", + "reason": f"[ERROR] LLM Call failed: {str(e)}", + "flag": "NEEDS_REVIEW" + }) -def classify_complaint(row: dict) -> dict: +def classify_complaint(row: dict, agents_text: str, skills_text: str) -> dict: """ Classify a single complaint row. Returns dict with: complaint_id, category, priority, reason, flag """ - raise NotImplementedError("Build this using your AI tool + agents.md") + description = row.get('description', '') + location = row.get('location', '') + + # Error handling: vague/short descriptions -> Other + NEEDS_REVIEW + if not description or len(description.strip()) < 5: + return { + "complaint_id": row.get('complaint_id', ''), + "category": "Other", + "priority": "Standard", + "reason": "Vague or short description", + "flag": "NEEDS_REVIEW" + } + + prompt = f"""You are the AI Complaint Classifier. + +=== AGENTS DEFINITION === +{agents_text} + +=== SKILLS DEFINITION === +{skills_text} + +=== INPUT COMPLAINT === +Location: {location} +Description: {description} + +=== INSTRUCTIONS === +Classify the complaint according to the schema and enforcement rules above. +Output strictly valid JSON with exactly the following keys: +"category", "priority", "reason", "flag" +""" + + response_text = call_llm(prompt) + + try: + result = json.loads(response_text) + result["complaint_id"] = row.get("complaint_id", "") + return result + except Exception as e: + return { + "complaint_id": row.get('complaint_id', ''), + "category": "Other", + "priority": "Standard", + "reason": f"Parsing failed: {str(e)}", + "flag": "NEEDS_REVIEW" + } def batch_classify(input_path: str, output_path: str): """Read input CSV, classify each row, write results CSV.""" - raise NotImplementedError("Build this using your AI tool + agents.md") + + base_dir = os.path.dirname(os.path.abspath(__file__)) + try: + with open(os.path.join(base_dir, 'agents.md'), 'r', encoding='utf-8') as f: + agents_text = f.read() + with open(os.path.join(base_dir, 'skills.md'), 'r', encoding='utf-8') as f: + skills_text = f.read() + except FileNotFoundError: + print("Error: agents.md or skills.md not found. Ensure they are in the same directory.") + return + + results = [] + + with open(input_path, mode='r', encoding='utf-8') as infile: + reader = csv.DictReader(infile) + for row in reader: + complaint_id = row.get('complaint_id') + # Error handling: malformed rows logged and skipped, processing continues + if not complaint_id: + print("Skipping malformed row missing complaint_id.") + continue + + print(f"Classifying {complaint_id}...") + result_dict = classify_complaint(row, agents_text, skills_text) + + clean_result = { + "complaint_id": result_dict.get("complaint_id", ""), + "category": result_dict.get("category", "Other"), + "priority": result_dict.get("priority", "Standard"), + "reason": result_dict.get("reason", "No reason provided"), + "flag": result_dict.get("flag", "") + } + results.append(clean_result) + + # Minimal sleep to avoid Gemini rate limits on the free tier + time.sleep(1.0) + + if not results: + print("No results to write.") + return + + with open(output_path, mode='w', encoding='utf-8', newline='') as outfile: + fieldnames = ["complaint_id", "category", "priority", "reason", "flag"] + writer = csv.DictWriter(outfile, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(results) if __name__ == "__main__": parser = argparse.ArgumentParser(description="UC-0A Complaint Classifier") diff --git a/uc-0a/results_pune.csv b/uc-0a/results_pune.csv new file mode 100644 index 0000000..ecbfa8a --- /dev/null +++ b/uc-0a/results_pune.csv @@ -0,0 +1,16 @@ +complaint_id,category,priority,reason,flag +PM-202401,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202402,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202406,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202408,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202410,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202411,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202413,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202418,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202419,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202420,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202427,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202428,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202430,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202433,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW +PM-202446,Other,Standard,[ERROR] GEMINI_API_KEY not set,NEEDS_REVIEW diff --git a/uc-0a/skills.md b/uc-0a/skills.md index 4e67823..5bc515a 100644 --- a/uc-0a/skills.md +++ b/uc-0a/skills.md @@ -1,15 +1,12 @@ -# skills.md — UC-0A Complaint Classifier -# INSTRUCTIONS: Same as agents.md — paste README into AI, ask for skills.md YAML - skills: - name: classify_complaint - description: "[FILL IN]" - input: "[FILL IN]" - output: "[FILL IN]" - error_handling: "[FILL IN]" + description: "Reads a single complaint and classifies it according to the predefined schema and enforcement rules, preventing taxonomy drift, severity blindness, missing justification, and hallucinated sub-categories." + input: "One complaint row (dict with description, location fields)" + output: "Dict with category, priority, reason, flag" + error_handling: "If the description is vague/short, or if ambiguity causes false confidence, it must fall back to outputting `category: Other` and `flag: NEEDS_REVIEW`." - name: batch_classify - description: "[FILL IN]" - input: "[FILL IN]" - output: "[FILL IN]" - error_handling: "[FILL IN]" + description: "Reads a CSV file containing multiple complaints, applies the classify_complaint skill to each row, and writes the results to a specified output CSV file." + input: "Path to test CSV file" + output: "Path to results CSV file" + error_handling: "If malformed rows are encountered, they must be logged and skipped, allowing processing of the remaining rows to continue." diff --git a/uc-mcp/agents.md b/uc-mcp/agents.md index d2e55c8..21984e1 100644 --- a/uc-mcp/agents.md +++ b/uc-mcp/agents.md @@ -1,32 +1,15 @@ -# agents.md — UC-MCP MCP Server -# INSTRUCTIONS: -# 1. Open your AI tool -# 2. Paste the full contents of uc-mcp/README.md -# 3. Use this prompt: -# "Read this UC README. Using the R.I.C.E framework, generate an -# agents.md YAML with four fields: role, intent, context, enforcement. -# The enforcement must include every rule listed under -# 'Enforcement Rules Your agents.md Must Include'. -# Output only valid YAML." -# 4. Paste the output below, replacing this placeholder -# 5. Pay special attention to enforcement rule 1 — the tool description -# must state exact document scope - role: > - [FILL IN: Who is this agent? What layer of the stack does it operate at? - Hint: an MCP server that exposes policy retrieval as a tool] + An MCP (Model Context Protocol) Server exposing a retrieval-augmented generation tool for CMC policy queries. intent: > - [FILL IN: What does a correctly implemented MCP server produce? - Hint: JSON-RPC compliant responses, scoped tool description, correct refusals] + To strictly define the scope of the `query_policy_documents` tool to prevent agents from hallucinating or answering out-of-scope questions, and to serve standard JSON-RPC 2.0 requests over HTTP. context: > - [FILL IN: What does this server have access to? - Hint: RAG server results only — no direct LLM calls, no outside knowledge] + The MCP server must interact exclusively through standard JSON-RPC requests for `tools/list` and `tools/call`, routing calls to the underlying RAG server. enforcement: - - "[FILL IN: Tool description scope rule]" - - "[FILL IN: Refusal documentation rule]" - - "[FILL IN: inputSchema required field rule]" - - "[FILL IN: isError on failure rule]" - - "[FILL IN: HTTP 200 for all JSON-RPC responses rule]" + - "Tool description must state the exact document scope: CMC HR Leave Policy, IT Acceptable Use Policy, Finance Reimbursement Policy." + - "Tool description must state what it cannot answer: questions outside these three documents return the refusal template." + - "inputSchema must require `question` as a non-empty string." + - "Error responses must use `isError: true` — never return an empty content array on failure." + - "The server must return HTTP 200 for all JSON-RPC responses including errors — transport errors use HTTP 4xx/5xx, application errors use JSON-RPC error objects." diff --git a/uc-mcp/mcp_server.py b/uc-mcp/mcp_server.py index 0400b6a..d6ce2d5 100644 --- a/uc-mcp/mcp_server.py +++ b/uc-mcp/mcp_server.py @@ -1,22 +1,3 @@ -""" -UC-MCP — mcp_server.py -Plain HTTP MCP Server — Starter File - -Build this using your AI coding tool: -1. Share agents.md, skills.md, and uc-mcp/README.md with your AI tool -2. Ask it to implement this file following the MCP protocol - described in the README -3. Run with: python3 mcp_server.py --port 8765 -4. Test with: python3 test_client.py --port 8765 - -Protocol: JSON-RPC 2.0 over HTTP POST -No external dependencies beyond Python stdlib. - -Methods to implement: - tools/list — return the tool definition for query_policy_documents - tools/call — execute query_policy_documents, return JSON-RPC response -""" - import json import argparse import sys @@ -37,20 +18,14 @@ # Import LLM adapter from llm_adapter import call_llm - # ── TOOL DEFINITION ────────────────────────────────────────────────────────── -# This is what the agent reads to decide when to call your tool. -# The description IS the enforcement — make it specific. TOOL_DEFINITION = { "name": "query_policy_documents", "description": ( - # FILL IN: Describe exactly what this tool covers and what it does not. - # Bad: "Answers questions about policies" - # Good: "Answers questions about CMC HR Leave Policy, IT Acceptable Use - # Policy, and Finance Reimbursement Policy only. Returns cited - # answers grounded in retrieved document chunks. Returns a refusal - # for questions outside these three documents." - "[FILL IN: specific scope + what it refuses]" + "Answers questions about CMC HR Leave Policy, IT Acceptable Use " + "Policy, and Finance Reimbursement Policy only. Returns cited " + "answers grounded in retrieved document chunks. Returns a refusal " + "for questions outside these three documents." ), "inputSchema": { "type": "object", @@ -64,49 +39,93 @@ }, } - # ── SKILL: query_policy_documents ──────────────────────────────────────────── def query_policy_documents(question: str) -> dict: """ Call the RAG server with the question. Return MCP content format: {"content": [...], "isError": bool} - - Error handling: - - If RAG refuses (no chunks above threshold) → isError: True - - If RAG raises exception → isError: True with error message """ - raise NotImplementedError( - "Implement query_policy_documents using your AI tool.\n" - "Hint: call rag_query(question, llm_call=call_llm), " - "check result['refused'], format as MCP content response." - ) - + try: + result = rag_query(question, llm_call=call_llm) + is_error = result.get('refused', False) + + content = [{ + "type": "text", + "text": result['answer'] + }] + + return { + "content": content, + "isError": is_error + } + except Exception as e: + return { + "content": [{ + "type": "text", + "text": f"Error calling RAG server: {str(e)}" + }], + "isError": True + } # ── SKILL: serve_mcp ───────────────────────────────────────────────────────── class MCPHandler(BaseHTTPRequestHandler): """ HTTP request handler implementing JSON-RPC 2.0. Handles POST requests to / with JSON-RPC body. - - Implement: - - tools/list → return TOOL_DEFINITION - - tools/call → call query_policy_documents, return result - - unknown methods → JSON-RPC error -32601 """ def do_POST(self): - raise NotImplementedError( - "Implement do_POST using your AI tool.\n" - "Hint: read Content-Length, parse JSON body, " - "dispatch on method, write JSON-RPC response.\n" - "Return HTTP 200 for all JSON-RPC responses including errors." - ) + content_length = int(self.headers.get('Content-Length', 0)) + if content_length == 0: + self.send_response(400) + self.end_headers() + return + + body = self.rfile.read(content_length) + + try: + req = json.loads(body) + except json.JSONDecodeError: + self.send_response(200) + self.send_header('Content-Type', 'application/json') + self.end_headers() + err_resp = {"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error"}, "id": None} + self.wfile.write(json.dumps(err_resp).encode('utf-8')) + return + + req_id = req.get('id') + method = req.get('method') + + response = {"jsonrpc": "2.0", "id": req_id} + + if method == 'tools/list': + response['result'] = {"tools": [TOOL_DEFINITION]} + elif method == 'tools/call': + params = req.get('params', {}) + tool_name = params.get('name') + + if tool_name == 'query_policy_documents': + args = params.get('arguments', {}) + question = args.get('question', '') + if not question: + response['error'] = {"code": -32602, "message": "Invalid params: question is required"} + else: + tool_result = query_policy_documents(question) + response['result'] = tool_result + else: + response['error'] = {"code": -32601, "message": "Tool not found"} + else: + response['error'] = {"code": -32601, "message": "Method not found"} + + self.send_response(200) + self.send_header('Content-Type', 'application/json') + self.end_headers() + self.wfile.write(json.dumps(response).encode('utf-8')) def log_message(self, format, *args): # Suppress default HTTP logging — use print for clarity print(f"[mcp_server] {args[0]} {args[1]}") - # ── MAIN ───────────────────────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser(description="UC-MCP Plain HTTP MCP Server") @@ -130,6 +149,5 @@ def main(): except KeyboardInterrupt: print("\n[mcp_server] Stopped.") - if __name__ == "__main__": main() diff --git a/uc-mcp/skills.md b/uc-mcp/skills.md index 5028507..37b7948 100644 --- a/uc-mcp/skills.md +++ b/uc-mcp/skills.md @@ -1,24 +1,12 @@ -# skills.md — UC-MCP MCP Server -# INSTRUCTIONS: -# 1. Open your AI tool -# 2. Paste the full contents of uc-mcp/README.md -# 3. Use this prompt: -# "Read this UC README. Generate a skills.md YAML defining the two -# skills: query_policy_documents and serve_mcp. Each skill needs: -# name, description, input, output, error_handling. -# error_handling must address the failure mode in the README. -# Output only valid YAML." -# 4. Paste the output below, replacing this placeholder - skills: - name: query_policy_documents - description: "[FILL IN]" - input: "[FILL IN: question string]" - output: "[FILL IN: MCP content format — content array + isError]" - error_handling: "[FILL IN: what happens when RAG refuses or raises exception]" + description: "Calls the underlying RAG server with a user question and returns the retrieved answer and cited sources." + input: "Takes: `question` (string)" + output: "Returns: answer + cited sources" + error_handling: "If RAG returns refused=True — return error content with `isError: true` and the refusal message." - name: serve_mcp - description: "[FILL IN]" - input: "[FILL IN: HTTP POST with JSON-RPC body]" - output: "[FILL IN: JSON-RPC 2.0 response, always HTTP 200]" - error_handling: "[FILL IN: unknown method → -32601, malformed request → -32700]" + description: "Starts a plain HTTP server that implements the MCP JSON-RPC protocol to expose tools to AI agents." + input: "Starts the HTTP server on a configurable port (default 8765) and Handles `tools/list` and `tools/call` requests" + output: "Returns JSON-RPC compliant responses" + error_handling: "unknown method → JSON-RPC error -32601" diff --git a/uc-rag/agents.md b/uc-rag/agents.md index 186c909..0fdf604 100644 --- a/uc-rag/agents.md +++ b/uc-rag/agents.md @@ -1,31 +1,15 @@ -# agents.md — UC-RAG RAG Server -# INSTRUCTIONS: -# 1. Open your AI tool -# 2. Paste the full contents of uc-rag/README.md -# 3. Use this prompt: -# "Read this UC README. Using the R.I.C.E framework, generate an -# agents.md YAML with four fields: role, intent, context, enforcement. -# Enforcement must include every rule listed under -# 'Enforcement Rules Your agents.md Must Include'. -# Output only valid YAML." -# 4. Paste the output below, replacing this placeholder -# 5. Check every enforcement rule against the README before saving - role: > - [FILL IN: Who is this agent? What is its operational boundary? - Hint: a retrieval-augmented policy assistant for city staff] + A retrieval-augmented generation (RAG) assistant for the City Municipal Corporation. It answers staff policy queries by retrieving relevant chunks from HR, IT, and Finance policy documents. intent: > - [FILL IN: What does a correct output look like? - Hint: answer + cited chunks + refusal when not covered] + To accurately answer staff policy questions by referencing only the retrieved document chunks, citing the sources, and explicitly refusing to answer questions not covered in the retrieved context. context: > - [FILL IN: What sources may the agent use? - Hint: retrieved chunks only — no general knowledge] + The assistant must strictly use only the provided chunks retrieved from the policy documents. It must exclude any external general knowledge or assumptions. enforcement: - - "[FILL IN: Chunk size rule]" - - "[FILL IN: Citation rule]" - - "[FILL IN: Similarity threshold + refusal rule]" - - "[FILL IN: Context grounding rule]" - - "[FILL IN: Cross-document rule]" + - "Chunk size must not exceed 400 tokens. Never split mid-sentence." + - "Every answer must cite the source document name and chunk index." + - "If no retrieved chunk scores above similarity threshold 0.6 — output the refusal template: 'This question is not covered in the retrieved policy documents. Retrieved chunks: [list chunk sources]. Please contact the relevant department for guidance.' Never generate an answer from general knowledge." + - "Answer must use only information present in the retrieved chunks. Never add context from outside the retrieved set." + - "If the query spans two documents — retrieve from each separately. Never merge retrieved chunks from different documents into one answer." diff --git a/uc-rag/rag_server.py b/uc-rag/rag_server.py index 3acfb1d..1fe1e96 100644 --- a/uc-rag/rag_server.py +++ b/uc-rag/rag_server.py @@ -1,23 +1,27 @@ -""" -UC-RAG — RAG Server -rag_server.py — Starter file - -Build this using your AI coding tool: -1. Share the contents of agents.md, skills.md, and uc-rag/README.md -2. Ask the AI to implement this file following the enforcement rules - in agents.md and the skill definitions in skills.md -3. Run with: python3 rag_server.py --build-index -4. Then: python3 rag_server.py --query "your question here" - -Stack: - pip3 install sentence-transformers chromadb - LLM: set your API key in llm_adapter.py (../uc-mcp/llm_adapter.py) - or set environment variable GEMINI_API_KEY -""" - import argparse import os import sys +import json +import re + +# Add path for llm_adapter +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'uc-mcp'))) +try: + from llm_adapter import call_llm +except ImportError: + print("Warning: llm_adapter not found. Please ensure it's at ../uc-mcp/llm_adapter.py") + def call_llm(prompt): + return "LLM call failed: llm_adapter not found." + +try: + from sentence_transformers import SentenceTransformer + import chromadb +except ImportError: + print("Warning: sentence_transformers or chromadb not installed.") + +# --- Helper --- +def token_count(text: str) -> int: + return len(text.split()) # --- SKILL: chunk_documents --- def chunk_documents(docs_dir: str, max_tokens: int = 400) -> list[dict]: @@ -25,24 +29,57 @@ def chunk_documents(docs_dir: str, max_tokens: int = 400) -> list[dict]: Load all .txt files from docs_dir. Split each into chunks of max_tokens, respecting sentence boundaries. Return list of: {doc_name, chunk_index, text} - - Failure mode to prevent: - - Never split mid-sentence (chunk boundary failure) - - Never exceed max_tokens per chunk """ - raise NotImplementedError( - "Implement chunk_documents using your AI tool.\n" - "Hint: use nltk.sent_tokenize or split on '. ' and accumulate " - "sentences until token limit is reached." - ) - + chunks = [] + + if not os.path.exists(docs_dir): + print(f"Error: {docs_dir} not found.") + return chunks + + for filename in os.listdir(docs_dir): + if not filename.endswith('.txt'): + continue + + filepath = os.path.join(docs_dir, filename) + with open(filepath, 'r', encoding='utf-8') as f: + content = f.read() + + sentences = re.split(r'(?<=[.!?])\s+', content) + + current_chunk_text = "" + chunk_idx = 0 + + for sentence in sentences: + sentence = sentence.strip() + if not sentence: continue + + if token_count(current_chunk_text) + token_count(sentence) <= max_tokens: + current_chunk_text += (sentence + " ") + else: + if current_chunk_text: + chunks.append({ + "doc_name": filename, + "chunk_index": chunk_idx, + "text": current_chunk_text.strip() + }) + chunk_idx += 1 + current_chunk_text = sentence + " " + + if current_chunk_text.strip(): + chunks.append({ + "doc_name": filename, + "chunk_index": chunk_idx, + "text": current_chunk_text.strip() + }) + + return chunks # --- SKILL: retrieve_and_answer --- def retrieve_and_answer( query: str, - collection, # ChromaDB collection - embedder, # SentenceTransformer model - llm_call, # callable: (prompt: str) -> str + collection, + embedder, + llm_call, top_k: int = 3, threshold: float = 0.6, ) -> dict: @@ -52,60 +89,144 @@ def retrieve_and_answer( If no chunks pass threshold, return refusal template. Otherwise call llm with retrieved chunks as context only. Return: {answer, cited_chunks: [{doc_name, chunk_index, score}]} - - Failure modes to prevent: - - Answer outside retrieved context - - Cross-document blending - - No citation """ - raise NotImplementedError( - "Implement retrieve_and_answer using your AI tool.\n" - "Hint: embed query, query ChromaDB collection, check distances, " - "build prompt with retrieved chunks only, call llm_call(prompt)." + query_emb = embedder.encode(query).tolist() + + results = collection.query( + query_embeddings=[query_emb], + n_results=top_k ) + + cited_chunks = [] + context_texts = [] + + # Chroma default space we used is cosine: distance = 1 - similarity + max_distance = 1.0 - threshold + + if results['distances'] and results['distances'][0]: + for i, dist in enumerate(results['distances'][0]): + if dist <= max_distance: + meta = results['metadatas'][0][i] + text = results['documents'][0][i] + cited_chunks.append({ + "doc_name": meta['doc_name'], + "chunk_index": meta['chunk_index'], + "score": 1.0 - dist + }) + context_texts.append(f"[{meta['doc_name']} Chunk {meta['chunk_index']}]: {text}") + + if not cited_chunks: + refusal = "This question is not covered in the retrieved policy documents. Retrieved chunks: None. Please contact the relevant department for guidance." + return {"answer": refusal, "cited_chunks": []} + + context_str = "\n\n".join(context_texts) + + base_dir = os.path.dirname(os.path.abspath(__file__)) + agents_path = os.path.join(base_dir, 'agents.md') + if os.path.exists(agents_path): + with open(agents_path, 'r', encoding='utf-8') as f: + agents_text = f.read() + else: + agents_text = "Enforce RICE rules." + + prompt = f"""You are the AI Assistant. + +=== AGENTS DEFINITION === +{agents_text} + +=== RETRIEVED CONTEXT === +{context_str} +=== USER QUERY === +{query} + +=== INSTRUCTIONS === +Answer the user query strictly using the retrieved context above. +Do not use any external knowledge. +You must cite the source document name and chunk index. +""" + + answer = llm_call(prompt) + + return { + "answer": answer, + "cited_chunks": cited_chunks + } # --- INDEX BUILDER --- def build_index(docs_dir: str, db_path: str = "./chroma_db"): - """ - Chunk all documents and store embeddings in ChromaDB. - Called once before querying. - """ - raise NotImplementedError( - "Implement build_index using your AI tool.\n" - "Hint: call chunk_documents(), embed each chunk with " - "SentenceTransformer, upsert into ChromaDB collection." + chunks = chunk_documents(docs_dir) + print(f"Generated {len(chunks)} chunks.") + + if not chunks: + print("No chunks to index.") + return + + print("Loading SentenceTransformer model...") + embedder = SentenceTransformer("all-MiniLM-L6-v2") + + print("Initializing ChromaDB...") + client = chromadb.PersistentClient(path=db_path) + + collection = client.get_or_create_collection( + name="policy_docs", + metadata={"hnsw:space": "cosine"} ) + + if collection.count() > 0: + collection.delete(ids=collection.get()['ids']) + + ids = [] + texts = [] + metadatas = [] + + for i, c in enumerate(chunks): + ids.append(f"chunk_{i}") + texts.append(c['text']) + metadatas.append({ + "doc_name": c['doc_name'], + "chunk_index": c['chunk_index'] + }) + + print(f"Embedding and upserting {len(texts)} chunks...") + embeddings = embedder.encode(texts).tolist() + + collection.upsert( + ids=ids, + embeddings=embeddings, + documents=texts, + metadatas=metadatas + ) + print("Successfully built the index.") - -# --- NAIVE MODE (run this first to see failure modes) --- +# --- NAIVE MODE --- def naive_query(query: str, docs_dir: str, llm_call): - """ - Load all documents into context without retrieval. - Run this BEFORE building your RAG pipeline to observe the failure modes. - """ - raise NotImplementedError( - "Implement naive_query using your AI tool.\n" - "Hint: load all .txt files, concatenate, pass to LLM with query. " - "No chunking, no retrieval, no enforcement." - ) + if not os.path.exists(docs_dir): + return f"Error: {docs_dir} not found." + + all_text = "" + for filename in os.listdir(docs_dir): + if filename.endswith('.txt'): + with open(os.path.join(docs_dir, filename), 'r', encoding='utf-8') as f: + all_text += f"\n--- {filename} ---\n{f.read()}" + + prompt = f"""You are a helpful assistant. +Context: +{all_text} +Query: {query} +Answer the query based on the context. +""" + return llm_call(prompt) # --- MAIN --- def main(): parser = argparse.ArgumentParser(description="UC-RAG RAG Server") - parser.add_argument("--build-index", action="store_true", - help="Build ChromaDB index from policy documents") - parser.add_argument("--query", type=str, - help="Query the RAG server") - parser.add_argument("--naive", action="store_true", - help="Run naive (no retrieval) mode to see failures") - parser.add_argument("--docs-dir", type=str, - default="../data/policy-documents", - help="Path to policy documents directory") - parser.add_argument("--db-path", type=str, - default="./chroma_db", - help="Path to ChromaDB storage directory") + parser.add_argument("--build-index", action="store_true", help="Build ChromaDB index") + parser.add_argument("--query", type=str, help="Query the RAG server") + parser.add_argument("--naive", action="store_true", help="Run naive mode") + parser.add_argument("--docs-dir", type=str, default="../data/policy-documents") + parser.add_argument("--db-path", type=str, default="./chroma_db") args = parser.parse_args() if not args.build_index and not args.query: @@ -118,18 +239,29 @@ def main(): print("Index built. Run with --query to test.") if args.query: + sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'uc-mcp'))) + from llm_adapter import call_llm + if args.naive: - # Import LLM adapter from uc-mcp - sys.path.insert(0, "../uc-mcp") - from llm_adapter import call_llm result = naive_query(args.query, args.docs_dir, call_llm) print(f"\nNaive answer:\n{result}") else: - # Full RAG query - raise NotImplementedError( - "Wire up retrieve_and_answer with ChromaDB and embedder here." + embedder = SentenceTransformer("all-MiniLM-L6-v2") + client = chromadb.PersistentClient(path=args.db_path) + collection = client.get_collection(name="policy_docs") + + result = retrieve_and_answer( + query=args.query, + collection=collection, + embedder=embedder, + llm_call=call_llm ) - + + print(f"\nAnswer:\n{result['answer']}\n") + if result['cited_chunks']: + print("Cited chunks:") + for c in result['cited_chunks']: + print(f"- {c['doc_name']} (Chunk {c['chunk_index']}) [Score: {c['score']:.2f}]") if __name__ == "__main__": main() diff --git a/uc-rag/skills.md b/uc-rag/skills.md index 167287b..baddbaa 100644 --- a/uc-rag/skills.md +++ b/uc-rag/skills.md @@ -1,25 +1,12 @@ -# skills.md — UC-RAG RAG Server -# INSTRUCTIONS: -# 1. Open your AI tool -# 2. Paste the full contents of uc-rag/README.md -# 3. Use this prompt: -# "Read this UC README. Generate a skills.md YAML defining the two -# skills: chunk_documents and retrieve_and_answer. Each skill needs: -# name, description, input, output, error_handling. -# error_handling must address the failure modes in the README. -# Output only valid YAML." -# 4. Paste the output below, replacing this placeholder -# 5. Verify error_handling addresses all three failure modes - skills: - name: chunk_documents - description: "[FILL IN]" - input: "[FILL IN: path to policy-documents directory]" - output: "[FILL IN: list of chunk dicts with doc_name, chunk_index, text]" - error_handling: "[FILL IN: what happens if a file is missing or unreadable]" + description: "Loads all policy documents from data/policy-documents/, splits them into chunks, and returns the chunks with metadata." + input: "None" + output: "List of chunks with metadata: {doc_name, chunk_index, text}" + error_handling: "Must strictly split on sentence boundaries and never exceed 400 tokens to prevent chunk boundary failures (e.g. splitting clauses)." - name: retrieve_and_answer - description: "[FILL IN]" - input: "[FILL IN: query string]" - output: "[FILL IN: answer string + list of cited chunks]" - error_handling: "[FILL IN: what happens when no chunk scores above 0.6]" + description: "Takes a query string, embeds it using sentence-transformers, retrieves relevant chunks from ChromaDB, and generates an answer using the LLM restricted to the retrieved context." + input: "Query string" + output: "Answer string + list of cited chunks" + error_handling: "Must filter out chunks scoring below 0.6 to prevent wrong chunk retrieval. If no chunk scores above 0.6, it must explicitly output the refusal template: 'This question is not covered in the retrieved policy documents. Retrieved chunks: [list chunk sources]. Please contact the relevant department for guidance.' to prevent answering outside the retrieved context." diff --git a/uc-rag/stub_rag.py b/uc-rag/stub_rag.py index 36fa00c..2eb5640 100644 --- a/uc-rag/stub_rag.py +++ b/uc-rag/stub_rag.py @@ -198,13 +198,31 @@ def retrieve_and_answer( f"[Source: {m['doc_name']}, chunk {m['chunk_index']}]\n{doc}" for doc, m, _ in passing ) + + base_dir = os.path.dirname(os.path.abspath(__file__)) + agents_path = os.path.join(base_dir, 'agents.md') + skills_path = os.path.join(base_dir, 'skills.md') + + agents_text = "" + if os.path.exists(agents_path): + with open(agents_path, 'r', encoding='utf-8') as f: + agents_text = f.read() + + skills_text = "" + if os.path.exists(skills_path): + with open(skills_path, 'r', encoding='utf-8') as f: + skills_text = f.read() + prompt = ( - f"Answer the following question using ONLY the provided context. " - f"Do not use any information outside the context. " - f"If the answer is not in the context, say so explicitly.\n\n" - f"Context:\n{context_blocks}\n\n" - f"Question: {query}\n\n" - f"Answer (cite source document and chunk for each claim):" + f"You are the AI Assistant.\n\n" + f"=== AGENTS DEFINITION ===\n{agents_text}\n\n" + f"=== SKILLS DEFINITION ===\n{skills_text}\n\n" + f"=== RETRIEVED CONTEXT ===\n{context_blocks}\n\n" + f"=== USER QUERY ===\n{query}\n\n" + f"=== INSTRUCTIONS ===\n" + f"Answer the user query strictly using the retrieved context above.\n" + f"Do not use any external knowledge.\n" + f"You must cite the source document name and chunk index.\n" ) if llm_call is None: