Spaces:
Sleeping
Sleeping
File size: 6,688 Bytes
2e818da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | 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]
]
|