diff --git a/docs/design.md b/docs/design.md new file mode 100644 index 0000000..c58065e --- /dev/null +++ b/docs/design.md @@ -0,0 +1,68 @@ +# ArgumentLab: Code Design & Implementation + +This document serves as a living guide to the actual codebase. It explains the current structure of the Python packages, what each file does, and how the logic is implemented. + +--- + +## Code Organization + +The application is modularized under `src/argument_lab/`, currently split into `core` data structures and the `orchestrator` graph logic. + +### 1. `core/models.py` (Pydantic Schemas) +This file defines the strict data schemas enforced throughout the system, primarily for structured LLM outputs. +- **`Argument`**: The core output format for agents. It strictly types the `agent` field (`"proponent"` or `"opponent"`) and enforces a `min_length=1` validator on `evidence`, meaning an agent cannot return an argument without at least one citation. +- **`EvidenceRef`** & **`Claim`**: Base objects for tracking retrieved evidence and registering claims. +- **`JudgeEvaluation` & `ArgumentScore`**: Defines the multi-dimensional scoring rubric (logical consistency, evidence support, relevance, completeness) as well as arrays for storing `hallucination_flags` and `contradiction_flags`. + +### 2. `core/state.py` (LangGraph State Management) +This file defines `DebateState`, the shared dictionary passed between all nodes in the LangGraph workflow. +Because nodes run in parallel, we implement custom **reducers** to prevent race conditions (where "last-write-wins" would corrupt the data): +- **`merge_dicts`**: Merges dictionary updates to `claims_registry` and `agent_positions`. +- **`union_sets`**: Merges sets of `addressed_claims` and `ignored_claims`. +- **`max_round`**: Ensures the `current_round` integer can only increase, preventing a lagging node from resetting the round number. +- **`merge_status`**: Resolves `status` updates by priority (e.g. if one node writes `"converged"` and another lazily writes `"in_progress"`, it resolves to `"converged"`). + +### 3. `core/agents.py` (Proponent & Opponent Logic) +Implements the core LangGraph agent nodes. Both agents follow a strict, deterministic pipeline instead of a chatty ReAct loop: +1. **Query Formulation**: A lightweight LLM call creates 1-3 targeted search queries based on the agent's stance and the debate history. +2. **Retrieval**: The queries are executed to fetch real `EvidenceRef` chunks. +3. **Generation**: The LLM uses `.with_structured_output(Argument)` to generate its argument, using the retrieved context. +4. **Counterpoint Enforcement**: In Rounds 2 and 3, the node explicitly re-prompts the LLM if it fails to populate `counterpoints_addressed` with an opponent's prior claim ID. + +### 4. `core/retriever.py` (RAG Interface) +A thin abstraction over the vector database (e.g. FAISS). It exposes `retrieve_multi()` which aggregates search results for multiple queries and deduplicates them by `source_id`, guaranteeing the best chunks are surfaced to the agent. + +### 5. `core/prompts.py` +Isolates all LangChain `ChatPromptTemplate` strings. It handles formatting debate histories and chunk excerpts, making it easy to iterate on prompt wording without touching workflow logic. + +### 6. `orchestrator/graph.py` (Workflow Topology) +This file compiles the `StateGraph` that controls the execution flow. It is heavily parallelized to reduce latency: +- **Agent Fan-out**: The `start_round` node branches unconditionally to `proponent_node` and `opponent_node`, running them concurrently. +- **Evaluation Sync & Fan-out**: Both agents join at a dummy node (`start_evaluation`). From there, the graph fans out again to three concurrent evaluation nodes: `judge_node`, `hallucination_check`, and `contradiction_check`. +- **Graph Update & Routing**: The parallel evaluation nodes join at `graph_update`, which writes final states. The `route_round` conditional edge then reads the state's `status` to decide whether to loop back to `start_round` or terminate the debate (`END`). + +## Implementation Efficiencies + +1. **The 2-step retrieval pipeline**: Doing `_formulate_queries` -> `_retrieve_evidence` before entering the structured argument generator avoids the grounding problem. It gives you the benefits of tool use without the risk of the LLM abandoning the schema or crashing into infinite tool loops. +2. **The `.model_copy(update=...)` filter**: + ```python + "evidence": [e for e in argument.evidence if e.source_id in valid_source_ids] or evidence_refs[:1] + ``` + If the LLM hallucinates source IDs, they are filtered out. But because Pydantic demands `min_length=1`, replacing an empty list with `evidence_refs[:1]` guarantees that validation will pass, avoiding a potential failure state. +3. **Counterpoint Enforcement**: Using the `_enforce_counterpoint_rule` to re-prompt the LLM explicitly when it fails to address an opponent's claim handles Option 3. Raising an `AgentError` if it fails twice propagates the failure to the workflow and it will get caught by the judge, as a result the judge will lower the score for logical consistency. +4. **State updates**: Extracting the confidence trajectory and accurately mapping `newly_ignored` claims via set math. + + +## Testing + +Here is what was added: +1. **`tests/core/test_state.py`**: Tests all the custom reducers (`union_sets`, `merge_dicts`, `max_round`, `merge_status`) to ensure they handle `None` defaults properly and execute the right merge logic. +2. **`tests/core/test_retriever.py`**: Mocks the `VectorIndex` protocol to test `Retriever.retrieve()` and ensures that `retrieve_multi()` correctly deduplicates source chunks, keeping the highest score. +3. **`tests/core/test_prompts.py`**: Tests the formatting helpers (`format_debate_history` and `format_evidence_context`) for edge cases like empty histories. +4. **`tests/core/test_agents.py`**: Tests the pure Python state derivation logic (`_get_prior_opponent_claim_ids` and `_update_state_from_argument`). +5. **`tests/orchestrator/test_graph.py`**: Tests that the `build_graph()` factory can successfully compile the graph topologically. +6. **`tests/core/test_models.py`**: Kept your existing test verifying `min_length=1` for evidence. + +--- + +*Note: This document should be updated whenever significant structural changes, new node implementations, or data models are introduced.* diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..b893048 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +pythonpath = src \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 7f7a083..0680ac2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,7 @@ uvicorn[standard]==0.30.6 # LLM / Agent frameworks langchain==0.2.14 +langchain-openai==0.1.22 langgraph==0.2.16 # OpenAI client (or compatible APIs) @@ -14,7 +15,7 @@ pydantic==2.9.2 pandas==2.2.3 # Vector / retrieval (optional but common) -faiss-cpu==1.8.0.post1 +faiss-cpu>=1.9.0.post1 # HTTP / integrations httpx==0.27.2 diff --git a/src/argument_lab/__init__.py b/src/argument_lab/__init__.py new file mode 100644 index 0000000..9cf13eb --- /dev/null +++ b/src/argument_lab/__init__.py @@ -0,0 +1 @@ +# Init \ No newline at end of file diff --git a/src/argument_lab/core/__init__.py b/src/argument_lab/core/__init__.py new file mode 100644 index 0000000..9cf13eb --- /dev/null +++ b/src/argument_lab/core/__init__.py @@ -0,0 +1 @@ +# Init \ No newline at end of file diff --git a/src/argument_lab/core/agents.py b/src/argument_lab/core/agents.py new file mode 100644 index 0000000..fefc672 --- /dev/null +++ b/src/argument_lab/core/agents.py @@ -0,0 +1,415 @@ +""" +argument_lab/core/agents.py + +Proponent and opponent agent nodes for the LangGraph debate workflow. + +Each agent follows a strict two-step pipeline: + Step 1 — Query formulation: a lightweight LLM call produces 1-3 search queries tailored to the agent's current goal. + Step 2 — Retrieval + generation: the queries are executed against the vector index, and the retrieved chunks are + injected into the system prompt before the structured argument generation call. + +This guarantees that every Argument object contains grounded evidence +before it ever reaches Pydantic validation. +""" + +import json +import uuid +from typing import Any + +from langchain_core.output_parsers import JsonOutputParser +from langchain_core.prompts import ChatPromptTemplate +from langchain_openai import ChatOpenAI + +from argument_lab.core.models import Argument, Claim, EvidenceRef +from argument_lab.core.retriever import Retriever, RetrieverError +from argument_lab.core.state import DebateState +from argument_lab.core.prompts import ( + QUERY_FORMULATION_SYSTEM, + QUERY_FORMULATION_USER, + AGENT_SYSTEM_TEMPLATE, + AGENT_USER_TEMPLATE, + COUNTERPOINT_RULES, + ROUND_GOALS, + format_debate_history, + format_evidence_context, +) + + +# --------------------------------------------------------------------------- +# LLM setup +# Temperature 0.2 for structured output — low enough for schema compliance, +# high enough to avoid degenerate repetition across rounds. +# --------------------------------------------------------------------------- + +import os + +_llm = ChatOpenAI(model="gpt-4o", temperature=0.2, api_key=os.environ.get("OPENAI_API_KEY", "dummy")) +_query_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.0, api_key=os.environ.get("OPENAI_API_KEY", "dummy")) + + + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +def _formulate_queries( + proposition: str, + stance: str, + history: str, + current_round: int, + llm: Any = _query_llm, +) -> list[str]: + """ + Step 1: Ask a lightweight LLM to produce search queries for this agent's + next argument. Returns a list of query strings, falling back to the + proposition itself if the LLM output cannot be parsed. + """ + prompt = ChatPromptTemplate.from_messages([ + ("system", QUERY_FORMULATION_SYSTEM), + ("user", QUERY_FORMULATION_USER), + ]) + chain = prompt | llm | JsonOutputParser() + + try: + result = chain.invoke({ + "proposition": proposition, + "stance": stance, + "history": history, + "current_round": current_round, + "round_goal": ROUND_GOALS[min(current_round, 3)], + }) + queries = result.get("queries", []) + if isinstance(queries, list) and all(isinstance(q, str) for q in queries): + return queries[:3] # enforce max + except Exception: + pass # fall through to fallback + + # Fallback: use the proposition directly so retrieval never returns empty + return [proposition] + + +def _retrieve_evidence( + retriever: Retriever, + queries: list[str], +) -> list[EvidenceRef]: + """ + Execute the queries against the vector index and convert chunks to EvidenceRef objects ready for the Argument schema. + Raises AgentError if retrieval returns nothing — no evidence means no valid argument can be produced. + """ + chunks = retriever.retrieve_multi(queries) + if not chunks: + raise AgentError( + "RAG retrieval returned no results. Cannot produce a grounded argument." + ) + return [ + EvidenceRef( + source_id=chunk.source_id, + excerpt=chunk.excerpt, + reliability_score=round(chunk.score, 3), + ) + for chunk in chunks + ] + + +def _generate_argument( + *, + role: str, + stance: str, + proposition: str, + current_round: int, + history: str, + evidence_refs: list[EvidenceRef], + evidence_context: str, + argument_id: str, + llm: Any = _llm, +) -> Argument: + """ + Step 2: Generate the structured Argument using the retrieved evidence injected into the system prompt. Uses .with_structured_output() to + enforce schema compliance at the LangChain layer. + """ + structured_llm = llm.with_structured_output(Argument) + + system_prompt = AGENT_SYSTEM_TEMPLATE.format_map({ + "role": role, + "stance": "FOR" if role == "Proponent" else "AGAINST", + "proposition": proposition, + "counterpoint_rule": COUNTERPOINT_RULES[min(current_round, 3)], + "evidence_context": evidence_context, + }) + + user_prompt = AGENT_USER_TEMPLATE.format_map({ + "history": history, + "current_round": current_round, + "argument_id": argument_id, + }) + + prompt = ChatPromptTemplate.from_messages([ + ("system", system_prompt), + ("user", user_prompt), + ]) + + chain = prompt | structured_llm + argument = chain.invoke({}) + + # Patch in the argument_id and round in case the LLM didn't follow them + # exactly — the schema enforces types but not specific string values. + argument = argument.model_copy(update={ + "id": argument_id, + "round": current_round, + "agent": role.lower(), + }) + + # Ensure the LLM cited sources that were actually retrieved (not hallucinated IDs) + valid_source_ids = {ref.source_id for ref in evidence_refs} + argument = argument.model_copy(update={ + "evidence": [e for e in argument.evidence if e.source_id in valid_source_ids] + or evidence_refs[:1], # guarantee min_length=1 even if LLM cited nothing valid + }) + + return argument + + +def _enforce_counterpoint_rule( + argument: Argument, + current_round: int, + prior_opponent_claim_ids: list[str], + role: str, + llm: Any = _llm, + proposition: str = "", + history: str = "", + evidence_refs: list[EvidenceRef] = [], + evidence_context: str = "", +) -> Argument: + """ + Option 3 enforcement: if Round >= 2 and counterpoints_addressed is empty, re-prompt the LLM once with an explicit correction instruction. + Raises AgentError if the second attempt also fails — this surfaces as a node failure rather than silently passing a non-compliant argument. + """ + if current_round < 2: + return argument # Round 1: no counterpoints required + + if argument.counterpoints_addressed: + return argument # already compliant + + if not prior_opponent_claim_ids: + # No opponent claims exist yet (e.g., opponent hasn't run this round) + # This shouldn't happen in normal flow but guard defensively. + return argument + + # Re-prompt with an explicit correction + correction_note = ( + f"Your previous response left counterpoints_addressed empty. " + f"You MUST include at least one of these opponent claim IDs: " + f"{prior_opponent_claim_ids}. Revise your argument to address " + f"at least one of these claims directly." + ) + structured_llm = llm.with_structured_output(Argument) + prompt = ChatPromptTemplate.from_messages([ + ("system", AGENT_SYSTEM_TEMPLATE.format_map({ + "role": role, + "stance": "FOR" if role == "Proponent" else "AGAINST", + "proposition": proposition, + "counterpoint_rule": COUNTERPOINT_RULES[min(current_round, 3)], + "evidence_context": evidence_context, + })), + ("user", AGENT_USER_TEMPLATE.format_map({ + "history": history, + "current_round": current_round, + "argument_id": argument.id, + })), + ("assistant", argument.model_dump_json()), + ("user", correction_note), + ]) + + revised = (prompt | structured_llm).invoke({}) + revised = revised.model_copy(update={ + "id": argument.id, + "round": current_round, + "agent": argument.agent, + }) + + if not revised.counterpoints_addressed: + raise AgentError( + f"{role} failed to address any opponent counterpoints in Round {current_round} " + f"after correction. Prior claim IDs available: {prior_opponent_claim_ids}" + ) + + return revised + + +def _get_prior_opponent_claim_ids(state: DebateState, my_role: str) -> list[str]: + """ + Returns claim IDs from the opponent's prior arguments. + Used to validate and enforce counterpoint_addressed in Round >= 2. + """ + opponent_role = "opponent" if my_role == "proponent" else "proponent" + return [ + arg.id + for arg in state.get("arguments", []) + if arg.agent == opponent_role and arg.round < state["current_round"] + ] + + +def _update_state_from_argument( + argument: Argument, + state: DebateState, +) -> dict: + """ + Derive all state updates that a new argument produces: + - adds argument to the accumulator + - registers its claim in claims_registry + - updates agent_positions with the new confidence score + - marks addressed and ignored claims + """ + # Register the new claim + from argument_lab.core.models import Claim + new_claim = Claim( + id=argument.id, + text=argument.claim, + agent=argument.agent, + round=argument.round, + ) + + # Determine which prior opponent claims were ignored this round + opponent_role = "opponent" if argument.agent == "proponent" else "proponent" + prior_opponent_ids = { + arg.id + for arg in state.get("arguments", []) + if arg.agent == opponent_role + } + newly_addressed = set(argument.counterpoints_addressed) + newly_ignored = prior_opponent_ids - newly_addressed - state.get("ignored_claims", set()) + + # Extend the agent's confidence trajectory + current_positions = dict(state.get("agent_positions", {})) + trajectory = list(current_positions.get(argument.agent, [])) + trajectory.append(argument.confidence_score) + + return { + "arguments": [argument], + "claims_registry": {argument.id: new_claim}, + "addressed_claims": newly_addressed, + "ignored_claims": newly_ignored, + "agent_positions": {argument.agent: trajectory}, + } + + +# --------------------------------------------------------------------------- +# Public node functions +# --------------------------------------------------------------------------- + +def make_proponent_node(retriever: Retriever): + """ + Factory that closes over a Retriever instance and returns a LangGraph- + compatible node function. Call this at graph compile time: + + workflow.add_node("proponent", make_proponent_node(retriever)) + """ + def proponent_node(state: DebateState) -> dict: + return _run_agent_node( + state=state, + role="Proponent", + agent_key="proponent", + retriever=retriever, + ) + return proponent_node + + +def make_opponent_node(retriever: Retriever): + """ + Factory that closes over a Retriever instance and returns a LangGraph- + compatible node function. + + workflow.add_node("opponent", make_opponent_node(retriever)) + """ + def opponent_node(state: DebateState) -> dict: + return _run_agent_node( + state=state, + role="Opponent", + agent_key="opponent", + retriever=retriever, + ) + return opponent_node + + +def _run_agent_node( + *, + state: DebateState, + role: str, + agent_key: str, + retriever: Retriever, +) -> dict: + """ + Shared implementation for both agent nodes. + + Pipeline: + 1. Format debate history for context + 2. Formulate search queries (lightweight LLM call) + 3. Retrieve evidence from vector index + 4. Generate structured Argument (main LLM call) + 5. Enforce counterpoint rule (re-prompt if needed) + 6. Derive and return state updates + """ + current_round = state["current_round"] + proposition = state["proposition"] + prior_arguments = state.get("arguments", []) + + # Step 1: format history + history = format_debate_history(prior_arguments) + stance = "FOR" if role == "Proponent" else "AGAINST" + + # Step 2: query formulation + queries = _formulate_queries( + proposition=proposition, + stance=stance, + history=history, + current_round=current_round, + ) + + # Step 3: retrieve evidence + evidence_refs = _retrieve_evidence(retriever, queries) + evidence_context = format_evidence_context( + # Pass the raw chunks back for display; evidence_refs are already converted + retriever.retrieve_multi(queries) + ) + + # Step 4: generate argument + argument_id = str(uuid.uuid4()) + argument = _generate_argument( + role=role, + stance=stance, + proposition=proposition, + current_round=current_round, + history=history, + evidence_refs=evidence_refs, + evidence_context=evidence_context, + argument_id=argument_id, + ) + + # Step 5: enforce counterpoint rule (Option 3) + prior_opponent_ids = _get_prior_opponent_claim_ids(state, agent_key) + argument = _enforce_counterpoint_rule( + argument=argument, + current_round=current_round, + prior_opponent_claim_ids=prior_opponent_ids, + role=role, + proposition=proposition, + history=history, + evidence_refs=evidence_refs, + evidence_context=evidence_context, + ) + + # Step 6: derive state updates + return _update_state_from_argument(argument, state) + + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + +class AgentError(RuntimeError): + """ + Raised when an agent node cannot produce a valid, schema-compliant argument. + The LangGraph node will propagate this as a node failure, which can be caught by a retry policy or surfaced to the metrics dashboard. + """ + pass \ No newline at end of file diff --git a/src/argument_lab/core/models.py b/src/argument_lab/core/models.py new file mode 100644 index 0000000..dc2ac46 --- /dev/null +++ b/src/argument_lab/core/models.py @@ -0,0 +1,44 @@ +from typing import Literal +from pydantic import BaseModel, Field + +class EvidenceRef(BaseModel): + source_id: str + excerpt: str + reliability_score: float = Field(ge=0.0, le=1.0) + +class Argument(BaseModel): + id: str + round: int + agent: Literal["proponent", "opponent"] + claim: str + evidence: list[EvidenceRef] = Field( + min_length=1, + description="Must contain \u22651 retrieved source" + ) + assumptions: list[str] + counterpoints_addressed: list[str] = Field( + default_factory=list, + description="Claim IDs of opponent's prior points" + ) + confidence_score: float = Field(ge=0.0, le=1.0) + +class Claim(BaseModel): + id: str + text: str + agent: str + round: int + +class ArgumentScore(BaseModel): + logical_consistency: float = Field(ge=0.0, le=1.0) + evidence_support: float = Field(ge=0.0, le=1.0) + relevance: float = Field(ge=0.0, le=1.0) + completeness: float = Field(ge=0.0, le=1.0) + +class JudgeEvaluation(BaseModel): + proponent_score: ArgumentScore + opponent_score: ArgumentScore + convergence_detected: bool + explanation: str + hallucination_flags: list[str] = Field(default_factory=list, + description="Claim IDs where evidence grounding failed") + contradiction_flags: list[str] = Field(default_factory=list) diff --git a/src/argument_lab/core/prompts.py b/src/argument_lab/core/prompts.py new file mode 100644 index 0000000..294d49d --- /dev/null +++ b/src/argument_lab/core/prompts.py @@ -0,0 +1,138 @@ +""" +argument_lab/core/prompts.py + +All prompt strings live here, outside the node logic. Keeping them +separate makes it easy to iterate on phrasing without touching orchestration +code, and makes prompt versioning straightforward. + +Templates use Python str.format_map() so they're readable without a +third-party templating library. +""" + +# --------------------------------------------------------------------------- +# Shared formatting helpers +# --------------------------------------------------------------------------- + +def format_argument(arg) -> str: + """Render a prior Argument object as a readable block for debate history.""" + evidence_lines = "\n".join( + f" [{e.source_id}] \"{e.excerpt}\" (reliability: {e.reliability_score:.2f})" + for e in arg.evidence + ) + addressed = ", ".join(arg.counterpoints_addressed) if arg.counterpoints_addressed else "none" + return ( + f"[{arg.agent.upper()} — Round {arg.round} — claim_id: {arg.id}]\n" + f"Claim: {arg.claim}\n" + f"Evidence:\n{evidence_lines}\n" + f"Assumptions: {', '.join(arg.assumptions) or 'none'}\n" + f"Counterpoints addressed: {addressed}\n" + f"Confidence: {arg.confidence_score:.2f}" + ) + + +def format_debate_history(arguments: list) -> str: + if not arguments: + return "No prior arguments." + return "\n\n".join(format_argument(a) for a in arguments) + + +def format_evidence_context(chunks: list) -> str: + """Render retrieved RAG chunks for injection into the generation prompt.""" + if not chunks: + return "No evidence retrieved." + return "\n".join( + f"[{c.source_id}] (similarity: {c.score:.2f})\n\"{c.excerpt}\"" + for c in chunks + ) + + +# --------------------------------------------------------------------------- +# Query formulation prompts +# Lightweight prompt used in Step 1 to get search queries from the LLM +# before the main argument generation call. +# --------------------------------------------------------------------------- + +QUERY_FORMULATION_SYSTEM = """\ +You are a research assistant for a structured debate. Your only job is to \ +formulate precise search queries that will retrieve the most relevant evidence \ +for the debater's next argument. + +Return a JSON object with a single key "queries" containing a list of 1-3 \ +short, specific search queries (each under 12 words). Do not explain. \ +Do not argue. Only return the JSON. +""" + +QUERY_FORMULATION_USER = """\ +Proposition under debate: {proposition} + +The debater you are helping argues: {stance} + +Debate history so far: +{history} + +Round {current_round} goal: {round_goal} + +Formulate search queries to find evidence for this debater's next argument. +""" + + +# --------------------------------------------------------------------------- +# Agent generation prompts +# Used in Step 2 after evidence has been retrieved and injected. +# --------------------------------------------------------------------------- + +AGENT_SYSTEM_TEMPLATE = """\ +You are the {role} in a structured multi-round debate. + +Your position: You argue {stance} the following proposition. +Proposition: "{proposition}" + +Rules of engagement: +1. Every claim you make MUST be grounded in the provided evidence. \ +Do not assert facts that are not present in the retrieved sources. +2. {counterpoint_rule} +3. Assign a confidence_score between 0.0 and 1.0 reflecting how strongly \ +the evidence supports your claim (not how strongly you personally believe it). +4. List any unstated premises your argument depends on in the assumptions field. +5. Your response must conform exactly to the required JSON schema. \ +No preamble. No explanation outside the schema. + +Retrieved evidence you MAY cite (you must cite at least one): +{evidence_context} +""" + +AGENT_USER_TEMPLATE = """\ +Debate history: +{history} + +Construct your Round {current_round} argument using the required schema. \ +Your argument id should be: "{argument_id}" +""" + + +# --------------------------------------------------------------------------- +# Round-specific rule strings (injected into AGENT_SYSTEM_TEMPLATE) +# --------------------------------------------------------------------------- + +COUNTERPOINT_RULES = { + 1: ( + "This is Round 1. No rebuttals are required. Focus on establishing " + "your strongest top-level case for your position. Leave counterpoints_addressed empty." + ), + 2: ( + "This is Round 2. You MUST address at least one specific claim from " + "your opponent's Round 1 argument. Include its claim_id in counterpoints_addressed. " + "Failing to address a prior claim will be penalized in scoring." + ), + 3: ( + "This is Round 3. You MUST address at least one claim from your opponent's " + "prior arguments. You may also update your confidence_score to reflect " + "any new evidence introduced in Round 2. Summarize your strongest remaining position." + ), +} + +ROUND_GOALS = { + 1: "Establish your strongest top-level case for your position.", + 2: "Rebut your opponent's Round 1 claims with specific evidence.", + 3: "Refine your position based on all prior evidence and finalize your case.", +} \ No newline at end of file diff --git a/src/argument_lab/core/retriever.py b/src/argument_lab/core/retriever.py new file mode 100644 index 0000000..6d5ea34 --- /dev/null +++ b/src/argument_lab/core/retriever.py @@ -0,0 +1,69 @@ +""" +argument_lab/core/retriever.py + +Thin abstraction over the vector index. Agents call retrieve() to get +grounded evidence before generating an argument. The implementation is +swappable (FAISS for MVP, OpenSearch for v2) — agents never import the +index directly. +""" + +from dataclasses import dataclass +from typing import Protocol + + +@dataclass +class RetrievedChunk: + source_id: str + excerpt: str + score: float # cosine similarity, [0, 1] + + +class VectorIndex(Protocol): + """ + Any object with this interface can be used as the backing index. + FAISS, ChromaDB, and OpenSearch all satisfy it with a thin wrapper. + """ + def similarity_search(self, query: str, k: int) -> list[RetrievedChunk]: + ... + + +class Retriever: + """ + Injected into each agent node at graph compile time via the config. + Agents call retrieve() with a natural-language query and get back + chunks they can directly attach as EvidenceRef objects. + """ + + def __init__(self, index: VectorIndex, top_k: int = 4): + self._index = index + self._top_k = top_k + + def retrieve(self, query: str) -> list[RetrievedChunk]: + """ + Returns up to top_k chunks ranked by similarity to the query. + Never raises — returns an empty list if the index is unavailable, + which the agent node treats as a hard failure (no evidence = no argument). + """ + try: + return self._index.similarity_search(query, k=self._top_k) + except Exception as exc: + # Propagate as a typed error so the node can surface it cleanly + raise RetrieverError(f"Index query failed: {exc}") from exc + + def retrieve_multi(self, queries: list[str]) -> list[RetrievedChunk]: + """ + Runs multiple queries and deduplicates by source_id, keeping the + highest-scoring chunk per source. Used when an agent formulates + separate queries for their main claim and their rebuttal. + """ + seen: dict[str, RetrievedChunk] = {} + for query in queries: + for chunk in self.retrieve(query): + existing = seen.get(chunk.source_id) + if existing is None or chunk.score > existing.score: + seen[chunk.source_id] = chunk + return sorted(seen.values(), key=lambda c: c.score, reverse=True) + + +class RetrieverError(RuntimeError): + pass \ No newline at end of file diff --git a/src/argument_lab/core/state.py b/src/argument_lab/core/state.py new file mode 100644 index 0000000..e203520 --- /dev/null +++ b/src/argument_lab/core/state.py @@ -0,0 +1,33 @@ +from typing import Annotated, Literal, TypedDict +import operator + +from argument_lab.core.models import Argument, Claim, JudgeEvaluation + +def union_sets(a: set[str] | None, b: set[str] | None) -> set[str]: + return (a or set()) | (b or set()) + +def merge_dicts(a: dict | None, b: dict | None) -> dict: + return {**(a or {}), **(b or {})} + +def max_round(a: int | None, b: int | None) -> int: + return max(a or 0, b or 0) + +def merge_status(a: str | None, b: str | None) -> str: + priority = ["terminated", "stalemate", "converged", "in_progress"] + a_val = a if a in priority else "in_progress" + b_val = b if b in priority else "in_progress" + return a_val if priority.index(a_val) < priority.index(b_val) else b_val + +class DebateState(TypedDict): + proposition: str + current_round: Annotated[int, max_round] + arguments: Annotated[list[Argument], operator.add] + claims_registry: Annotated[dict[str, Claim], merge_dicts] + addressed_claims: Annotated[set[str], union_sets] + ignored_claims: Annotated[set[str], union_sets] + agent_positions: Annotated[dict[str, list[float]], merge_dicts] + repetition_flags: Annotated[list[str], operator.add] + status: Annotated[Literal["in_progress", "converged", "stalemate", "terminated"], merge_status] + hallucination_flags: Annotated[list[str], operator.add] + contradiction_flags: Annotated[list[str], operator.add] + scores: Annotated[list[JudgeEvaluation], operator.add] diff --git a/src/argument_lab/orchestrator/__init__.py b/src/argument_lab/orchestrator/__init__.py new file mode 100644 index 0000000..9cf13eb --- /dev/null +++ b/src/argument_lab/orchestrator/__init__.py @@ -0,0 +1 @@ +# Init \ No newline at end of file diff --git a/src/argument_lab/orchestrator/graph.py b/src/argument_lab/orchestrator/graph.py new file mode 100644 index 0000000..b90a767 --- /dev/null +++ b/src/argument_lab/orchestrator/graph.py @@ -0,0 +1,107 @@ +""" +argument_lab/core/graph.py + +Builds and compiles the LangGraph debate workflow. + +Agent nodes are constructed via factories (make_proponent_node, +make_opponent_node) that close over a shared Retriever instance. +Pass a configured Retriever to build_graph() at startup. +""" + +from langgraph.graph import StateGraph, START, END + +from argument_lab.core.agents import make_proponent_node, make_opponent_node +from argument_lab.core.retriever import Retriever +from argument_lab.core.state import DebateState + + +def start_round(state: DebateState) -> dict: + return {} + + +def judge_node(state: DebateState) -> dict: + return {} + + +def hallucination_check(state: DebateState) -> dict: + return {} + + +def contradiction_check(state: DebateState) -> dict: + return {} + + +def graph_update(state: DebateState) -> dict: + return {} + + +def route_round(state: DebateState) -> str: + if state.get("status") in ["converged", "stalemate", "terminated"]: + return END + if state.get("current_round", 1) > 3: + return END + return "start_round" + + +def build_graph(retriever: Retriever): + """ + Compile the debate workflow. Call once at application startup and + reuse the compiled graph across debate sessions. + + Usage: + retriever = Retriever(index=your_faiss_index) + debate_graph = build_graph(retriever) + result = debate_graph.invoke({ + "proposition": "...", + "current_round": 1, + "arguments": [], + "claims_registry": {}, + "addressed_claims": set(), + "ignored_claims": set(), + "agent_positions": {}, + "repetition_flags": [], + "status": "in_progress", + "hallucination_flags": [], + "contradiction_flags": [], + "scores": [], + }) + """ + workflow = StateGraph(DebateState) + + # --- Node registration --- + workflow.add_node("start_round", start_round) + workflow.add_node("proponent", make_proponent_node(retriever)) + workflow.add_node("opponent", make_opponent_node(retriever)) + workflow.add_node("start_evaluation", lambda state: {}) + workflow.add_node("judge", judge_node) + workflow.add_node("hallucination_check", hallucination_check) + workflow.add_node("contradiction_check", contradiction_check) + workflow.add_node("graph_update", graph_update) + + # --- Edge wiring --- + + # Entry + workflow.add_edge(START, "start_round") + + # Fan-out: both agents run in parallel each round + workflow.add_edge("start_round", "proponent") + workflow.add_edge("start_round", "opponent") + + # Fan-in: both agents must complete before evaluation starts + workflow.add_edge("proponent", "start_evaluation") + workflow.add_edge("opponent", "start_evaluation") + + # Fan-out: judge, hallucination, and contradiction run in parallel + workflow.add_edge("start_evaluation", "judge") + workflow.add_edge("start_evaluation", "hallucination_check") + workflow.add_edge("start_evaluation", "contradiction_check") + + # Fan-in: all evaluation nodes complete before graph_update + workflow.add_edge("judge", "graph_update") + workflow.add_edge("hallucination_check", "graph_update") + workflow.add_edge("contradiction_check", "graph_update") + + # Conditional routing: continue or terminate + workflow.add_conditional_edges("graph_update", route_round) + + return workflow.compile() \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..9cf13eb --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +# Init \ No newline at end of file diff --git a/tests/core/__init__.py b/tests/core/__init__.py new file mode 100644 index 0000000..9cf13eb --- /dev/null +++ b/tests/core/__init__.py @@ -0,0 +1 @@ +# Init \ No newline at end of file diff --git a/tests/core/test_agents.py b/tests/core/test_agents.py new file mode 100644 index 0000000..dea140c --- /dev/null +++ b/tests/core/test_agents.py @@ -0,0 +1,39 @@ +import pytest +from argument_lab.core.agents import _get_prior_opponent_claim_ids, _update_state_from_argument +from argument_lab.core.models import Argument + +def test_get_prior_opponent_claim_ids(): + state = { + "current_round": 2, + "arguments": [ + Argument.model_construct(id="a1", round=1, agent="proponent", claim="p1"), + Argument.model_construct(id="a2", round=1, agent="opponent", claim="o1"), + ] + } + ids = _get_prior_opponent_claim_ids(state, "proponent") + assert ids == ["a2"] + + ids = _get_prior_opponent_claim_ids(state, "opponent") + assert ids == ["a1"] + +def test_update_state_from_argument(): + arg = Argument.model_construct( + id="a2", round=2, agent="proponent", claim="p2", + counterpoints_addressed=["o1"], confidence_score=0.9 + ) + + state = { + "arguments": [ + Argument.model_construct(id="o1", round=1, agent="opponent"), + Argument.model_construct(id="o2", round=1, agent="opponent") + ], + "ignored_claims": set(), + "agent_positions": {"proponent": [0.8]} + } + + updates = _update_state_from_argument(arg, state) + + assert "a2" in updates["claims_registry"] + assert updates["addressed_claims"] == {"o1"} + assert updates["ignored_claims"] == {"o2"} + assert updates["agent_positions"]["proponent"] == [0.8, 0.9] diff --git a/tests/core/test_models.py b/tests/core/test_models.py new file mode 100644 index 0000000..d52624e --- /dev/null +++ b/tests/core/test_models.py @@ -0,0 +1,31 @@ +import pytest +from pydantic import ValidationError +from argument_lab.core.models import Argument, EvidenceRef + +def test_argument_requires_evidence(): + with pytest.raises(ValidationError) as exc_info: + Argument( + id="arg_1", + round=1, + agent="proponent", + claim="AI is good.", + evidence=[], + assumptions=[], + counterpoints_addressed=[], + confidence_score=0.9 + ) + assert "at least 1" in str(exc_info.value).lower() or "min_length" in str(exc_info.value).lower() + +def test_valid_argument(): + ev = EvidenceRef(source_id="doc_1", excerpt="AI helps.", reliability_score=0.8) + arg = Argument( + id="arg_1", + round=1, + agent="proponent", + claim="AI is good.", + evidence=[ev], + assumptions=[], + counterpoints_addressed=[], + confidence_score=0.9 + ) + assert arg.evidence[0].source_id == "doc_1" diff --git a/tests/core/test_prompts.py b/tests/core/test_prompts.py new file mode 100644 index 0000000..d52624e --- /dev/null +++ b/tests/core/test_prompts.py @@ -0,0 +1,31 @@ +import pytest +from pydantic import ValidationError +from argument_lab.core.models import Argument, EvidenceRef + +def test_argument_requires_evidence(): + with pytest.raises(ValidationError) as exc_info: + Argument( + id="arg_1", + round=1, + agent="proponent", + claim="AI is good.", + evidence=[], + assumptions=[], + counterpoints_addressed=[], + confidence_score=0.9 + ) + assert "at least 1" in str(exc_info.value).lower() or "min_length" in str(exc_info.value).lower() + +def test_valid_argument(): + ev = EvidenceRef(source_id="doc_1", excerpt="AI helps.", reliability_score=0.8) + arg = Argument( + id="arg_1", + round=1, + agent="proponent", + claim="AI is good.", + evidence=[ev], + assumptions=[], + counterpoints_addressed=[], + confidence_score=0.9 + ) + assert arg.evidence[0].source_id == "doc_1" diff --git a/tests/core/test_retriever.py b/tests/core/test_retriever.py new file mode 100644 index 0000000..e50a8cc --- /dev/null +++ b/tests/core/test_retriever.py @@ -0,0 +1,39 @@ +import pytest +from argument_lab.core.retriever import Retriever, RetrievedChunk, RetrieverError + +class MockIndex: + def similarity_search(self, query: str, k: int) -> list[RetrievedChunk]: + if query == "fail": + raise RuntimeError("Index failure") + return [ + RetrievedChunk(source_id="doc_1", excerpt=f"match for {query}", score=0.9), + RetrievedChunk(source_id="doc_2", excerpt="other", score=0.5) + ] + +def test_retrieve(): + retriever = Retriever(index=MockIndex(), top_k=2) + chunks = retriever.retrieve("test") + assert len(chunks) == 2 + assert chunks[0].source_id == "doc_1" + +def test_retrieve_failure(): + retriever = Retriever(index=MockIndex(), top_k=2) + with pytest.raises(RetrieverError): + retriever.retrieve("fail") + +def test_retrieve_multi_dedup(): + class MockDedupIndex: + def similarity_search(self, query: str, k: int) -> list[RetrievedChunk]: + if query == "q1": + return [RetrievedChunk("doc_1", "foo", 0.9)] + if query == "q2": + return [RetrievedChunk("doc_1", "foo", 0.95), RetrievedChunk("doc_2", "bar", 0.8)] + return [] + + retriever = Retriever(index=MockDedupIndex(), top_k=2) + chunks = retriever.retrieve_multi(["q1", "q2"]) + + assert len(chunks) == 2 + assert chunks[0].source_id == "doc_1" + assert chunks[0].score == 0.95 + assert chunks[1].source_id == "doc_2" diff --git a/tests/core/test_state.py b/tests/core/test_state.py new file mode 100644 index 0000000..955000a --- /dev/null +++ b/tests/core/test_state.py @@ -0,0 +1,24 @@ +import pytest +from argument_lab.core.state import union_sets, merge_dicts, max_round, merge_status + +def test_union_sets(): + assert union_sets({"a"}, {"b"}) == {"a", "b"} + assert union_sets(None, {"b"}) == {"b"} + assert union_sets({"a"}, None) == {"a"} + assert union_sets(None, None) == set() + +def test_merge_dicts(): + assert merge_dicts({"a": 1}, {"b": 2}) == {"a": 1, "b": 2} + assert merge_dicts({"a": 1}, {"a": 2}) == {"a": 2} # right hand wins + assert merge_dicts(None, {"b": 2}) == {"b": 2} + +def test_max_round(): + assert max_round(1, 2) == 2 + assert max_round(3, 1) == 3 + assert max_round(None, 2) == 2 + +def test_merge_status(): + assert merge_status("in_progress", "converged") == "converged" + assert merge_status("stalemate", "converged") == "stalemate" + assert merge_status("terminated", "in_progress") == "terminated" + assert merge_status(None, "converged") == "converged" diff --git a/tests/orchestrator/__init__.py b/tests/orchestrator/__init__.py new file mode 100644 index 0000000..9cf13eb --- /dev/null +++ b/tests/orchestrator/__init__.py @@ -0,0 +1 @@ +# Init \ No newline at end of file diff --git a/tests/orchestrator/test_graph.py b/tests/orchestrator/test_graph.py new file mode 100644 index 0000000..090a5b8 --- /dev/null +++ b/tests/orchestrator/test_graph.py @@ -0,0 +1,13 @@ +import pytest +from argument_lab.orchestrator.graph import build_graph +from argument_lab.core.retriever import Retriever, RetrievedChunk + +class DummyIndex: + def similarity_search(self, query: str, k: int) -> list[RetrievedChunk]: + return [RetrievedChunk("d1", "ex", 0.9)] + +def test_build_graph_compiles(): + retriever = Retriever(DummyIndex()) + graph = build_graph(retriever) + + assert graph is not None