Spaces:
Sleeping
Sleeping
| import json | |
| import os | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Annotated, Any, TypedDict | |
| import faiss | |
| import numpy as np | |
| import streamlit as st | |
| from duckduckgo_search import DDGS | |
| from langchain_core.messages import AIMessage, HumanMessage, SystemMessage | |
| from langchain_groq import ChatGroq | |
| from langgraph.graph import END, StateGraph | |
| from langgraph.graph.message import add_messages | |
| from sentence_transformers import SentenceTransformer | |
| # --- Settings --- | |
| GROQ_API_KEY = os.getenv("GROQ_API_KEY", "") | |
| GROQ_MODEL = os.getenv("GROQ_MODEL", "openai/gpt-oss-120b") | |
| EMBEDDING_MODEL = "all-MiniLM-L6-v2" | |
| MAX_RESEARCH_ITERATIONS = 2 | |
| MEMORY_DIR = Path(__file__).parent / "data" / "memory" | |
| MEMORY_DIR.mkdir(parents=True, exist_ok=True) | |
| FAISS_INDEX_PATH = MEMORY_DIR / "research_memory.index" | |
| FAISS_META_PATH = MEMORY_DIR / "research_memory_meta.json" | |
| AGENT_COLORS = { | |
| "Researcher": "#4A90D9", | |
| "Critic": "#E67E22", | |
| "Synthesizer": "#9B59B6", | |
| "Writer": "#27AE60", | |
| } | |
| # --- Memory --- | |
| class ResearchMemory: | |
| def __init__(self) -> None: | |
| self._encoder = SentenceTransformer(EMBEDDING_MODEL) | |
| self._dim = self._encoder.get_sentence_embedding_dimension() | |
| self._index: faiss.IndexFlatL2 | None = None | |
| self._metadata: list[dict[str, Any]] = [] | |
| self._load() | |
| def _load(self) -> None: | |
| if FAISS_INDEX_PATH.exists() and FAISS_META_PATH.exists(): | |
| self._index = faiss.read_index(str(FAISS_INDEX_PATH)) | |
| self._metadata = json.loads(FAISS_META_PATH.read_text(encoding="utf-8")) | |
| else: | |
| self._index = faiss.IndexFlatL2(self._dim) | |
| self._metadata = [] | |
| def _save(self) -> None: | |
| if self._index is not None: | |
| faiss.write_index(self._index, str(FAISS_INDEX_PATH)) | |
| FAISS_META_PATH.write_text( | |
| json.dumps(self._metadata, indent=2, ensure_ascii=False), | |
| encoding="utf-8", | |
| ) | |
| def _embed(self, text: str) -> np.ndarray: | |
| vec = self._encoder.encode([text], normalize_embeddings=True) | |
| return np.array(vec, dtype=np.float32) | |
| def store(self, topic: str, report: str) -> None: | |
| entry = { | |
| "topic": topic, | |
| "report": report, | |
| "timestamp": datetime.now(timezone.utc).isoformat(), | |
| } | |
| vector = self._embed(f"{topic}\n{report[:2000]}") | |
| self._index.add(vector) | |
| self._metadata.append(entry) | |
| self._save() | |
| def recall(self, query: str, k: int = 3) -> list[dict[str, Any]]: | |
| if not self._metadata or self._index is None or self._index.ntotal == 0: | |
| return [] | |
| vector = self._embed(query) | |
| distances, indices = self._index.search(vector, min(k, self._index.ntotal)) | |
| results = [] | |
| for dist, idx in zip(distances[0], indices[0]): | |
| if idx < 0 or idx >= len(self._metadata): | |
| continue | |
| entry = dict(self._metadata[idx]) | |
| entry["score"] = float(dist) | |
| results.append(entry) | |
| return results | |
| def list_sessions(self) -> list[dict[str, Any]]: | |
| return [ | |
| { | |
| "topic": m["topic"], | |
| "timestamp": m["timestamp"], | |
| "preview": m["report"][:200] + "..." if len(m["report"]) > 200 else m["report"], | |
| } | |
| for m in reversed(self._metadata) | |
| ] | |
| def clear(self) -> None: | |
| self._index = faiss.IndexFlatL2(self._dim) | |
| self._metadata = [] | |
| self._save() | |
| # --- Search --- | |
| def web_search(query: str, max_results: int = 5) -> list[dict[str, str]]: | |
| try: | |
| with DDGS() as ddgs: | |
| results = list(ddgs.text(query, max_results=max_results)) | |
| return [ | |
| {"title": r.get("title", ""), "snippet": r.get("body", ""), "url": r.get("href", "")} | |
| for r in results | |
| ] | |
| except Exception: | |
| return [] | |
| def format_search_results(results: list[dict[str, str]]) -> str: | |
| if not results: | |
| return "No web results found. Use your training knowledge and note assumptions." | |
| lines = [] | |
| for i, r in enumerate(results, 1): | |
| lines.append(f"[{i}] {r['title']}\n URL: {r['url']}\n {r['snippet']}") | |
| return "\n\n".join(lines) | |
| # --- LangGraph state & agents --- | |
| class ResearchState(TypedDict): | |
| topic: str | |
| memory_context: str | |
| research_notes: str | |
| sources: list[str] | |
| critique: str | |
| critique_approved: bool | |
| synthesis: str | |
| report: str | |
| iteration: int | |
| messages: Annotated[list, add_messages] | |
| def _llm() -> ChatGroq: | |
| return ChatGroq(groq_api_key=GROQ_API_KEY, model_name=GROQ_MODEL, temperature=0.3) | |
| def _invoke(system: str, user: str) -> str: | |
| response = _llm().invoke([SystemMessage(content=system), HumanMessage(content=user)]) | |
| return response.content.strip() | |
| def researcher_node(state: ResearchState) -> dict: | |
| topic = state["topic"] | |
| search_results = web_search(f"{topic} research overview latest developments", max_results=6) | |
| sources = [r["url"] for r in search_results if r.get("url")] | |
| user_parts = [f"Topic: {topic}"] | |
| if state.get("memory_context"): | |
| user_parts.append(f"Relevant past research from memory:\n{state['memory_context']}") | |
| if state.get("research_notes") and state.get("critique"): | |
| user_parts.append( | |
| f"Previous research (needs improvement):\n{state['research_notes']}\n\n" | |
| f"Critic feedback to address:\n{state['critique']}" | |
| ) | |
| user_parts.append(f"Web search results:\n{format_search_results(search_results)}") | |
| notes = _invoke( | |
| "You are a Research Agent. Gather comprehensive, factual research on the given topic. " | |
| "Structure your notes with: Key Concepts, Current State, Important Developments, " | |
| "Open Questions, and Data Points. Be thorough and cite sources by URL when available.", | |
| "\n\n".join(user_parts), | |
| ) | |
| return { | |
| "research_notes": notes, | |
| "sources": sources, | |
| "messages": [AIMessage(content=f"[Researcher] Completed research on: {topic}")], | |
| } | |
| def critic_node(state: ResearchState) -> dict: | |
| iteration = state.get("iteration", 0) | |
| critique = _invoke( | |
| "You are a Critic Agent. Rigorously evaluate research quality. " | |
| "Check for: factual gaps, weak reasoning, missing perspectives, outdated info, " | |
| "and lack of evidence. End your response with exactly one line: " | |
| "VERDICT: APPROVED or VERDICT: REVISE", | |
| f"Topic: {state['topic']}\n\nResearch notes to evaluate:\n{state['research_notes']}", | |
| ) | |
| approved = "VERDICT: APPROVED" in critique.upper() | |
| force_approve = iteration >= MAX_RESEARCH_ITERATIONS - 1 | |
| return { | |
| "critique": critique, | |
| "critique_approved": approved or force_approve, | |
| "iteration": iteration + 1, | |
| "messages": [ | |
| AIMessage( | |
| content=f"[Critic] {'Approved' if (approved or force_approve) else 'Requested revision'}" | |
| ) | |
| ], | |
| } | |
| def synthesizer_node(state: ResearchState) -> dict: | |
| user_parts = [ | |
| f"Topic: {state['topic']}", | |
| f"Research notes:\n{state['research_notes']}", | |
| f"Critic evaluation:\n{state['critique']}", | |
| ] | |
| if state.get("memory_context"): | |
| user_parts.append(f"Historical context from memory:\n{state['memory_context']}") | |
| synthesis = _invoke( | |
| "You are a Synthesizer Agent. Merge research findings and critique into a unified " | |
| "analytical framework. Identify themes, connections, contradictions, and insights. " | |
| "Produce a structured synthesis ready for report writing.", | |
| "\n\n".join(user_parts), | |
| ) | |
| return { | |
| "synthesis": synthesis, | |
| "messages": [AIMessage(content="[Synthesizer] Unified analytical framework ready")], | |
| } | |
| def writer_node(state: ResearchState) -> dict: | |
| sources = state.get("sources", []) | |
| source_block = "\n".join(f"- {url}" for url in sources) if sources else "No external URLs captured." | |
| report = _invoke( | |
| "You are a Writer Agent. Produce a polished, publication-quality deep-research report. " | |
| "Use markdown formatting with: Executive Summary, Introduction, Key Findings, " | |
| "Deep Analysis, Implications, Limitations, and Conclusion. " | |
| "Write in clear, authoritative prose. Include a Sources section at the end.", | |
| f"Topic: {state['topic']}\n\nSynthesis to expand into full report:\n{state['synthesis']}\n\n" | |
| f"Available sources:\n{source_block}", | |
| ) | |
| return { | |
| "report": report, | |
| "messages": [AIMessage(content="[Writer] Deep-research report completed")], | |
| } | |
| def _route_after_critic(state: ResearchState) -> str: | |
| return "synthesizer" if state.get("critique_approved", False) else "researcher" | |
| def build_research_graph(): | |
| graph = StateGraph(ResearchState) | |
| graph.add_node("researcher", researcher_node) | |
| graph.add_node("critic", critic_node) | |
| graph.add_node("synthesizer", synthesizer_node) | |
| graph.add_node("writer", writer_node) | |
| graph.set_entry_point("researcher") | |
| graph.add_edge("researcher", "critic") | |
| graph.add_conditional_edges( | |
| "critic", | |
| _route_after_critic, | |
| {"researcher": "researcher", "synthesizer": "synthesizer"}, | |
| ) | |
| graph.add_edge("synthesizer", "writer") | |
| graph.add_edge("writer", END) | |
| return graph.compile() | |
| # --- Streamlit UI --- | |
| st.set_page_config(page_title="Deep Research Agents", page_icon="🔬", layout="wide") | |
| def init_session(): | |
| if "memory" not in st.session_state: | |
| st.session_state.memory = ResearchMemory() | |
| if "graph" not in st.session_state: | |
| st.session_state.graph = build_research_graph() | |
| if "running" not in st.session_state: | |
| st.session_state.running = False | |
| def render_agent_status(agent: str, status: str): | |
| color = AGENT_COLORS.get(agent, "#888") | |
| st.markdown( | |
| f'<div style="padding:8px 12px;margin:4px 0;border-left:4px solid {color};' | |
| f'background:#1a1a2e;border-radius:4px;">' | |
| f'<strong style="color:{color}">{agent}</strong> — {status}</div>', | |
| unsafe_allow_html=True, | |
| ) | |
| def run_research(topic: str, use_memory: bool): | |
| memory: ResearchMemory = st.session_state.memory | |
| graph = st.session_state.graph | |
| memory_context = "" | |
| if use_memory: | |
| recalled = memory.recall(topic, k=3) | |
| if recalled: | |
| memory_context = "\n\n---\n\n".join( | |
| f"Past topic: {r['topic']} (relevance score: {r['score']:.3f})\n{r['report'][:1500]}" | |
| for r in recalled | |
| ) | |
| progress = st.progress(0, text="Initializing agents...") | |
| status_box = st.empty() | |
| stages = [ | |
| ("Researcher", "Gathering information..."), | |
| ("Critic", "Evaluating research quality..."), | |
| ("Synthesizer", "Building analytical framework..."), | |
| ("Writer", "Composing final report..."), | |
| ] | |
| initial_state = { | |
| "topic": topic, | |
| "memory_context": memory_context, | |
| "research_notes": "", | |
| "sources": [], | |
| "critique": "", | |
| "critique_approved": False, | |
| "synthesis": "", | |
| "report": "", | |
| "iteration": 0, | |
| "messages": [], | |
| } | |
| final_state = None | |
| stage_idx = 0 | |
| for event in graph.stream(initial_state, stream_mode="updates"): | |
| for node_name, node_output in event.items(): | |
| stage_idx = {"researcher": 0, "critic": 1, "synthesizer": 2, "writer": 3}.get(node_name, stage_idx) | |
| progress.progress(min((stage_idx + 1) / len(stages), 1.0), text=stages[stage_idx][1]) | |
| with status_box.container(): | |
| for i, (agent, desc) in enumerate(stages): | |
| if i <= stage_idx: | |
| render_agent_status(agent, "✓ " + desc if i < stage_idx else desc) | |
| if final_state is None: | |
| final_state = {**initial_state, **node_output} | |
| else: | |
| final_state.update(node_output) | |
| progress.progress(1.0, text="Research complete!") | |
| return final_state | |
| def main(): | |
| init_session() | |
| memory: ResearchMemory = st.session_state.memory | |
| st.title("🔬 Deep Research Agents") | |
| st.caption( | |
| "Multi-agent research pipeline — **Researcher** → **Critic** → **Synthesizer** → **Writer** " | |
| "powered by Groq + LangGraph + FAISS" | |
| ) | |
| with st.sidebar: | |
| st.header("Settings") | |
| use_memory = st.toggle("Use cross-session memory", value=True) | |
| st.divider() | |
| st.subheader("Past Sessions") | |
| sessions = memory.list_sessions() | |
| if sessions: | |
| for s in sessions[:10]: | |
| st.markdown(f"**{s['topic']}**") | |
| st.caption(s["timestamp"][:19]) | |
| st.text(s["preview"][:120]) | |
| st.divider() | |
| else: | |
| st.info("No past sessions yet.") | |
| if st.button("Clear memory", type="secondary"): | |
| memory.clear() | |
| st.rerun() | |
| if not GROQ_API_KEY: | |
| st.error( | |
| "Set your `GROQ_API_KEY` environment variable. " | |
| "Get a free key at [console.groq.com](https://console.groq.com)." | |
| ) | |
| st.stop() | |
| col1, col2 = st.columns([3, 1]) | |
| with col1: | |
| topic = st.text_input( | |
| "Research topic", | |
| placeholder="e.g. The impact of quantum computing on cryptography", | |
| ) | |
| with col2: | |
| st.write("") | |
| st.write("") | |
| submit = st.button("Start Research", type="primary", disabled=st.session_state.running) | |
| if submit and topic.strip(): | |
| st.session_state.running = True | |
| with st.spinner("Agents collaborating..."): | |
| result = run_research(topic.strip(), use_memory) | |
| if result and result.get("report"): | |
| st.success("Report generated successfully!") | |
| tabs = st.tabs(["📄 Report", "🔍 Research Notes", "⚖️ Critique", "🧩 Synthesis"]) | |
| with tabs[0]: | |
| st.markdown(result["report"]) | |
| with tabs[1]: | |
| st.markdown(result.get("research_notes", "")) | |
| with tabs[2]: | |
| st.markdown(result.get("critique", "")) | |
| with tabs[3]: | |
| st.markdown(result.get("synthesis", "")) | |
| if use_memory: | |
| memory.store(topic.strip(), result["report"]) | |
| st.toast("Saved to cross-session memory", icon="💾") | |
| st.download_button( | |
| "Download Report", | |
| data=result["report"], | |
| file_name=f"research_{topic[:40].replace(' ', '_')}.md", | |
| mime="text/markdown", | |
| ) | |
| else: | |
| st.error("Research pipeline did not produce a report. Check your API key and try again.") | |
| st.session_state.running = False | |
| with st.expander("How it works"): | |
| st.markdown( | |
| """ | |
| | Agent | Role | | |
| |-------|------| | |
| | **Researcher** | Web search + knowledge synthesis into structured notes | | |
| | **Critic** | Quality gate — gaps, bias, evidence checks; loops back if needed | | |
| | **Synthesizer** | Merges research + critique into unified analytical framework | | |
| | **Writer** | Produces publication-quality markdown report | | |
| """ | |
| ) | |
| if __name__ == "__main__": | |
| main() | |