python -m venv venv
source venv/bin/activate
pip install httpx openai arxiv pydantic
export OPENAI_API_KEY="sk-your-key"

import arxiv
import httpx
import logging
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


class ArxivClient:
    """Search and retrieve papers from arxiv."""

    def search(self, query: str, max_results: int = 20, sort_by: str = "relevance") -> List[Dict[str, Any]]:
        sort_map = {
            "relevance": arxiv.SortCriterion.Relevance,
            "date": arxiv.SortCriterion.SubmittedDate,
        }
        client = arxiv.Client()
        search = arxiv.Search(
            query=query,
            max_results=max_results,
            sort_by=sort_map.get(sort_by, arxiv.SortCriterion.Relevance),
        )
        papers = []
        for result in client.results(search):
            papers.append({
                "id": result.entry_id,
                "title": result.title,
                "authors": [a.name for a in result.authors],
                "abstract": result.summary,
                "published": result.published.isoformat(),
                "updated": result.updated.isoformat(),
                "pdf_url": result.pdf_url,
                "categories": result.categories,
                "primary_category": result.primary_category,
                "source": "arxiv",
            })
        return papers


class SemanticScholarClient:
    """Search and retrieve papers from Semantic Scholar."""

    BASE_URL = "https://api.semanticscholar.org/graph/v1"

    def __init__(self, api_key: Optional[str] = None):
        self.headers = {}
        if api_key:
            self.headers["x-api-key"] = api_key

    async def search(self, query: str, limit: int = 20) -> List[Dict[str, Any]]:
        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"{self.BASE_URL}/paper/search",
                params={"query": query, "limit": limit, "fields": "title,authors,abstract,year,citationCount,url"},
                headers=self.headers,
            )
            data = response.json()
        return [
            {
                "id": p.get("paperId", ""),
                "title": p.get("title", ""),
                "authors": [a.get("name", "") for a in p.get("authors", [])],
                "abstract": p.get("abstract", ""),
                "year": p.get("year"),
                "citations": p.get("citationCount", 0),
                "url": p.get("url", ""),
                "source": "semantic_scholar",
            }
            for p in data.get("data", [])
        ]

    async def get_paper(self, paper_id: str) -> Dict[str, Any]:
        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"{self.BASE_URL}/paper/{paper_id}",
                params={"fields": "title,authors,abstract,year,citationCount,references,tldr"},
                headers=self.headers,
            )
            return response.json()

from openai import AsyncOpenAI
import json


class PaperExtractor:
    """Extract structured information from research papers using LLM."""

    def __init__(self, model: str = "gpt-4o"):
        self.client = AsyncOpenAI()
        self.model = model

    async def extract(self, paper: Dict[str, Any]) -> Dict[str, Any]:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {
                    "role": "system",
                    "content": """Extract structured information from this research paper.
                    Return JSON:
                    {
                        "research_question": "main question addressed",
                        "methodology": "approach used",
                        "key_findings": ["list of findings"],
                        "contributions": ["list of contributions"],
                        "limitations": ["list of limitations"],
                        "future_work": ["suggested future directions"],
                        "key_metrics": {"metric_name": "value"},
                        "dataset_used": "dataset name or description"
                    }""",
                },
                {"role": "user", "content": f"Title: {paper['title']}\n\nAbstract: {paper.get('abstract', '')}"},
            ],
            temperature=0.0,
        )
        try:
            return json.loads(response.choices[0].message.content)
        except (json.JSONDecodeError, IndexError):
            return {"research_question": "Unknown", "key_findings": []}

    async def extract_batch(self, papers: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        import asyncio
        tasks = [self.extract(paper) for paper in papers]
        results = await asyncio.gather(*tasks)
        return [{**paper, "extracted": ext} for paper, ext in zip(papers, results)]

class PaperSynthesizer:
    """Synthesize findings across multiple research papers."""

    def __init__(self, model: str = "gpt-4o"):
        self.client = AsyncOpenAI()
        self.model = model

    async def synthesize(self, papers: List[Dict[str, Any]], topic: str) -> Dict[str, Any]:
        summaries = []
        for p in papers:
            ext = p.get("extracted", {})
            summaries.append(f"Paper: {p['title']} ({p.get('year', 'N/A')})")
            summaries.append(f"  Method: {ext.get('methodology', 'N/A')}")
            summaries.append(f"  Findings: {', '.join(ext.get('key_findings', [])[:3])}")
        papers_text = "\n".join(summaries)
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {
                    "role": "system",
                    "content": """Synthesize findings across multiple papers.
                    Identify: consensus, contradictions, methodological trends, and overall narrative.
                    Use academic writing style with citations (Author, Year).""",
                },
                {"role": "user", "content": f"Topic: {topic}\n\nPapers:\n{papers_text}\n\nSynthesize the literature:"},
            ],
            temperature=0.3,
        )
        return {"topic": topic, "num_papers": len(papers), "synthesis": response.choices[0].message.content}


class GapAnalyzer:
    """Identify research gaps and future directions."""

    def __init__(self, model: str = "gpt-4o"):
        self.client = AsyncOpenAI()
        self.model = model

    async def analyze_gaps(self, papers: List[Dict[str, Any]], topic: str) -> Dict[str, Any]:
        findings = []
        for p in papers:
            ext = p.get("extracted", {})
            findings.append(f"- {p['title']}: {', '.join(ext.get('key_findings', [])[:2])}")
        findings_text = "\n".join(findings)
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {
                    "role": "system",
                    "content": """Analyze research gaps in this literature.
                    Return JSON:
                    {
                        "gaps": [{"gap": "description", "importance": "high|medium|low", "suggested_approach": "how to address"}],
                        "underexplored_areas": ["list of areas needing more research"],
                        "methodological_gaps": ["methods not yet applied"],
                        "future_directions": ["promising research directions"]
                    }""",
                },
                {"role": "user", "content": f"Topic: {topic}\n\nFindings:\n{findings_text}"},
            ],
            temperature=0.3,
        )
        try:
            return json.loads(response.choices[0].message.content)
        except (json.JSONDecodeError, IndexError):
            return {"gaps": [], "future_directions": []}


class LiteratureReviewAgent:
    """Orchestrate the complete literature review pipeline."""

    def __init__(self, model: str = "gpt-4o"):
        self.arxiv = ArxivClient()
        self.s2 = SemanticScholarClient()
        self.extractor = PaperExtractor(model)
        self.synthesizer = PaperSynthesizer(model)
        self.gap_analyzer = GapAnalyzer(model)

    async def review(self, topic: str, max_papers: int = 15) -> Dict[str, Any]:
        arxiv_papers = self.arxiv.search(topic, max_results=max_papers)
        s2_papers = await self.s2.search(topic, limit=max_papers)
        all_papers = self._deduplicate(arxiv_papers + s2_papers)
        extracted = await self.extractor.extract_batch(all_papers[:max_papers])
        synthesis = await self.synthesizer.synthesize(extracted, topic)
        gaps = await self.gap_analyzer.analyze_gaps(extracted, topic)
        return {
            "topic": topic,
            "total_papers": len(extracted),
            "papers": extracted,
            "synthesis": synthesis,
            "research_gaps": gaps,
        }

    def _deduplicate(self, papers: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        seen = set()
        unique = []
        for p in papers:
            title = p.get("title", "").lower().strip()
            if title not in seen:
                seen.add(title)
                unique.append(p)
        return unique

    def generate_review_report(self, review_data: Dict[str, Any]) -> str:
        report = f"# Literature Review: {review_data['topic']}\n\n"
        report += f"## Overview\nReviewed {review_data['total_papers']} papers\n\n"
        report += f"## Synthesis\n{review_data['synthesis']['synthesis']}\n\n"
        report += "## Key Papers\n"
        for p in review_data["papers"][:10]:
            report += f"- {p['title']} ({p.get('year', 'N/A')})\n"
        report += "\n## Research Gaps\n"
        for gap in review_data.get("research_gaps", {}).get("gaps", []):
            report += f"- {gap.get('gap', '')} (Importance: {gap.get('importance', 'medium')})\n"
        return report