Spaces:
Sleeping
Sleeping
| import asyncio | |
| from typing import Any, Dict, List, Literal, Optional | |
| from pydantic import BaseModel, Field | |
| from app.agents.cerebras_client import CerebrasClient | |
| from app.rag.chromadb_client import ChromaDBClient | |
| from app.rag.models import ConsumerType, EvidenceRequest | |
| from app.rag.retrieval_service import PaperEvidenceService | |
| from app.rag.vector_indexes import PaperEvidenceIndex | |
| class _ClaimItem(BaseModel): | |
| claim: str = Field(description="A single factual claim asserted in the draft") | |
| citation: str = Field( | |
| default="", | |
| description="The citation marker attached to the claim (e.g. '[1]', 'Smith 2020'), or '' if none", | |
| ) | |
| class _ClaimExtraction(BaseModel): | |
| claims: List[_ClaimItem] = Field( | |
| default_factory=list, | |
| description="1-6 key factual claims from the draft that carry a citation", | |
| ) | |
| class _Verdict(BaseModel): | |
| verdict: Literal["supported", "partial", "unsupported", "contradicted", "not-found"] | |
| evidence: str = Field( | |
| description="One sentence grounded in the source passages that justifies the verdict " | |
| "(or notes their absence). Do not invent facts beyond the passages." | |
| ) | |
| class DraftAgent: | |
| """Verifies the cited claims in a student's draft against the project's own | |
| uploaded corpus (ChromaDB), one claim at a time, grounded -> never guesses. | |
| """ | |
| def __init__(self, client: Optional[CerebrasClient] = None, db: Optional[ChromaDBClient] = None) -> None: | |
| self._client = client or CerebrasClient() | |
| self._db = db | |
| def _get_db(self) -> ChromaDBClient: | |
| if self._db is None: | |
| self._db = ChromaDBClient() | |
| return self._db | |
| async def verify_draft(self, project_id: str, draft_text: str) -> List[Dict[str, Any]]: | |
| loop = asyncio.get_event_loop() | |
| # 1. Extract cited claims from the draft. | |
| try: | |
| extraction = await loop.run_in_executor(None, lambda: self._client.structured_complete( | |
| [ | |
| { | |
| "role": "system", | |
| "content": ( | |
| "Extract 1-6 key factual claims from the student's draft that are attributed " | |
| "to a source (a citation marker like [1], (Smith, 2020), etc.). Return each " | |
| "claim as a self-contained sentence plus the citation marker it carries. " | |
| "Ignore rhetorical or unsupported sentences with no citation." | |
| ), | |
| }, | |
| {"role": "user", "content": draft_text[:8000]}, | |
| ], | |
| _ClaimExtraction, | |
| )) | |
| claims = extraction.claims | |
| except Exception as e: | |
| print("Draft claim extraction error:", e) | |
| claims = [] | |
| if not claims: | |
| return [{ | |
| "claim": "No cited claims found in the draft.", | |
| "verdict": "not-found", | |
| "evidence": "Add citations (e.g. [1]) next to your factual statements so they can be verified.", | |
| "citation": "N/A", | |
| }] | |
| # 2. Verify each claim against the project corpus, concurrently. | |
| return await asyncio.gather(*(self._verify_claim(project_id, c) for c in claims)) | |
| async def _verify_claim(self, project_id: str, claim: _ClaimItem) -> Dict[str, Any]: | |
| loop = asyncio.get_event_loop() | |
| citation = claim.citation or "N/A" | |
| chunks = await loop.run_in_executor(None, lambda: self._retrieve(project_id, claim.claim)) | |
| if not chunks: | |
| return { | |
| "claim": claim.claim, | |
| "verdict": "not-found", | |
| "evidence": "No matching material was found in this project's uploaded corpus.", | |
| "citation": citation, | |
| } | |
| ctx = "\n\n".join(f"[{c.get('source', '?')}] {c['text']}" for c in chunks) | |
| try: | |
| verdict = await loop.run_in_executor(None, lambda: self._client.structured_complete( | |
| [ | |
| { | |
| "role": "system", | |
| "content": ( | |
| "You verify a single claim against retrieved source passages from the " | |
| "student's own corpus. Classify strictly from the passages:\n" | |
| "- supported: the passages directly affirm the claim.\n" | |
| "- partial: the passages affirm part of it or affirm it with caveats.\n" | |
| "- unsupported: the passages are on-topic but do not establish the claim.\n" | |
| "- contradicted: the passages assert the opposite.\n" | |
| "- not-found: the passages are unrelated to the claim.\n" | |
| "Ground your one-sentence evidence in the passages -> never invent facts." | |
| ), | |
| }, | |
| {"role": "user", "content": f"CLAIM: {claim.claim}\n\nSOURCE PASSAGES:\n{ctx}"}, | |
| ], | |
| _Verdict, | |
| )) | |
| return { | |
| "claim": claim.claim, | |
| "verdict": verdict.verdict, | |
| "evidence": verdict.evidence, | |
| "citation": citation, | |
| } | |
| except Exception as e: | |
| print("Draft verify error:", e) | |
| return { | |
| "claim": claim.claim, | |
| "verdict": "not-found", | |
| "evidence": "Verification could not be completed for this claim.", | |
| "citation": citation, | |
| } | |
| def _retrieve(self, project_id: str, query: str, n: int = 4) -> List[Dict[str, Any]]: | |
| service = PaperEvidenceService( | |
| index=None if self._db is None else PaperEvidenceIndex(self._get_db()) | |
| ) | |
| result = service.retrieve(EvidenceRequest( | |
| query=query, | |
| consumer=ConsumerType.DRAFT, | |
| project_id=project_id, | |
| token_budget=max(512, n * 500), | |
| project_memory_policy="none", | |
| student_memory_policy="none", | |
| )) | |
| citations = {citation.evidence_id: citation for citation in result.citations} | |
| return [ | |
| { | |
| "text": item.evidence.index_text, | |
| "source": citations.get(item.evidence.evidence_id).filename | |
| if item.evidence.evidence_id in citations else item.evidence.document_id, | |
| "evidence_id": item.evidence.evidence_id, | |
| "page_start": item.evidence.page_start, | |
| } | |
| for item in result.evidence[:n] | |
| ] | |