AI_tools

  1from abc import ABC, abstractmethod
  2from dataclasses import dataclass
  3import os
  4from typing import cast
  5
  6from dotenv import load_dotenv
  7from google import genai
  8from google.genai import types
  9from openai import OpenAI
 10from pydantic import BaseModel
 11
 12from config import ENV_PATH
 13
 14_ = load_dotenv(ENV_PATH)
 15
 16
 17@dataclass
 18class PaperContent:
 19    """
 20    Wraps available content for a paper.
 21    Functions use the richest source available: pdf > full_text > abstract.
 22    Populate full_text (TeX) or pdf (bytes) via downloads.py when needed.
 23    """
 24    abstract: str
 25    full_text: str | None = None  # raw TeX source
 26    pdf: bytes | None = None      # PDF bytes
 27
 28    def best_text(self) -> str:
 29        if self.full_text:
 30            return self.full_text
 31        return self.abstract
 32
 33
 34# Response schemas
 35
 36class _TagResponse(BaseModel):
 37    tags: list[str]
 38
 39
 40class SummaryResult(BaseModel):
 41    tldr: str
 42    key_contributions: list[str]
 43
 44
 45class _RelatedResponse(BaseModel):
 46    related_ids: list[str]
 47
 48
 49# ---------------------------------------------------------------------------
 50# AIProvider ABC
 51# ---------------------------------------------------------------------------
 52
 53class AIProvider(ABC):
 54    """Unified interface for AI-powered paper analysis."""
 55
 56    @abstractmethod
 57    def tag(self, content: PaperContent) -> list[str]:
 58        """Generate 3-5 relevant tags for a paper."""
 59
 60    @abstractmethod
 61    def summarize(self, content: PaperContent) -> SummaryResult:
 62        """Return a one-sentence TLDR and 2-4 key contributions."""
 63
 64    @abstractmethod
 65    def find_related(
 66        self,
 67        content: PaperContent,
 68        candidates: list[tuple[str, str]],
 69        threshold: int = 5,
 70    ) -> list[str]:
 71        """Return IDs of the most conceptually related papers from candidates."""
 72
 73
 74# ---------------------------------------------------------------------------
 75# GeminiProvider
 76# ---------------------------------------------------------------------------
 77
 78class GeminiProvider(AIProvider):
 79    """AI provider backed by Google Gemini."""
 80
 81    def __init__(self) -> None:
 82        self._client = None  # lazy init
 83
 84    def _get_client(self):
 85        if self._client is None:
 86            api_key = os.getenv("GENAI_API_KEY_TAG_GEN")
 87            if not api_key:
 88                raise EnvironmentError("GENAI_API_KEY_TAG_GEN not set.")
 89            self._client = genai.Client(api_key=api_key)
 90        return self._client
 91
 92    def _generate(self, prompt: str, content: PaperContent, schema: type[BaseModel]) -> BaseModel:
 93        parts: list = [types.Part.from_text(text=prompt)]
 94        if content.pdf:
 95            parts.append(types.Part.from_bytes(data=content.pdf, mime_type="application/pdf"))
 96        elif content.full_text:
 97            parts.append(types.Part.from_text(text=content.full_text))
 98        else:
 99            parts.append(types.Part.from_text(text=content.abstract))
100        response = self._get_client().models.generate_content(
101            model="gemini-2.0-flash",
102            contents=parts,
103            config=types.GenerateContentConfig(
104                response_mime_type="application/json",
105                response_schema=schema,
106            ),
107        )
108        return cast(BaseModel, response.parsed)
109
110    def tag(self, content: PaperContent) -> list[str]:
111        parsed = cast(_TagResponse, self._generate(
112            "Generate 3-5 relevant Obsidian tags for this academic paper.",
113            content, _TagResponse,
114        ))
115        return [f"#{t.strip().lstrip('#').replace(' ', '_')}" for t in parsed.tags]
116
117    def summarize(self, content: PaperContent) -> SummaryResult:
118        return cast(SummaryResult, self._generate(
119            "Summarize this academic paper into a one-sentence TLDR and 2-4 key contributions.",
120            content, SummaryResult,
121        ))
122
123    def find_related(
124        self,
125        content: PaperContent,
126        candidates: list[tuple[str, str]],
127        threshold: int = 5,
128    ) -> list[str]:
129        candidate_block = "\n\n".join(
130            f"ID: {pid}\n{ab}" for pid, ab in candidates[:40])
131        parsed = cast(_RelatedResponse, self._generate(
132            f"Which of the following papers are most conceptually related to this one? "
133            f"Return up to {threshold} paper IDs.\n\n{candidate_block}",
134            content, _RelatedResponse,
135        ))
136        return parsed.related_ids
137
138
139# ---------------------------------------------------------------------------
140# OpenAIProvider
141# ---------------------------------------------------------------------------
142
143class OpenAIProvider(AIProvider):
144    """AI provider backed by OpenAI (GPT-4o, etc.)."""
145
146    def __init__(self, model: str = "gpt-4o-mini") -> None:
147        self._client = None  # lazy init
148        self._model = model
149
150    def _get_client(self):
151        if self._client is None:
152            api_key = os.getenv("OPENAI_API_KEY")
153            if not api_key:
154                raise EnvironmentError("OPENAI_API_KEY not set.")
155            self._client = OpenAI(api_key=api_key)
156        return self._client
157
158    def _generate(self, prompt: str, content: PaperContent, schema: type[BaseModel]) -> BaseModel:
159        text = content.best_text()
160        response = self._get_client().beta.chat.completions.parse(
161            model=self._model,
162            messages=[
163                {"role": "system", "content": prompt},
164                {"role": "user", "content": text},
165            ],
166            response_format=schema,
167        )
168        result = response.choices[0].message.parsed
169        assert result is not None
170        return result
171
172    def tag(self, content: PaperContent) -> list[str]:
173        parsed = cast(_TagResponse, self._generate(
174            "Generate 3-5 relevant Obsidian tags for this academic paper. "
175            "Return as JSON with a 'tags' array of strings.",
176            content, _TagResponse,
177        ))
178        return [f"#{t.strip().lstrip('#').replace(' ', '_')}" for t in parsed.tags]
179
180    def summarize(self, content: PaperContent) -> SummaryResult:
181        return cast(SummaryResult, self._generate(
182            "Summarize this academic paper into a one-sentence TLDR and 2-4 key contributions. "
183            "Return as JSON with 'tldr' (string) and 'key_contributions' (array of strings).",
184            content, SummaryResult,
185        ))
186
187    def find_related(
188        self,
189        content: PaperContent,
190        candidates: list[tuple[str, str]],
191        threshold: int = 5,
192    ) -> list[str]:
193        candidate_block = "\n\n".join(
194            f"ID: {pid}\n{ab}" for pid, ab in candidates[:40])
195        parsed = cast(_RelatedResponse, self._generate(
196            f"Which of the following papers are most conceptually related to this one? "
197            f"Return up to {threshold} paper IDs as JSON with a 'related_ids' array.\n\n{candidate_block}",
198            content, _RelatedResponse,
199        ))
200        return parsed.related_ids
201
202
203# ---------------------------------------------------------------------------
204# Module-level active provider + public API
205# ---------------------------------------------------------------------------
206
207_provider: AIProvider | None = None
208
209
210def _get_provider() -> AIProvider:
211    global _provider
212    if _provider is None:
213        _provider = GeminiProvider()
214    return _provider
215
216
217def set_provider(provider: AIProvider) -> None:
218    """Switch the active AI provider."""
219    global _provider
220    _provider = provider
221
222
223def tag(content: PaperContent, file_path: str | None = None) -> list[str]:
224    """Generate 3-5 Obsidian tags. Optionally append to file_path."""
225    tags = _get_provider().tag(content)
226    if file_path:
227        with open(file_path, "a", encoding="utf-8") as f:
228            f.write("\n" + " ".join(tags))
229    return tags
230
231
232def summarize(content: PaperContent) -> SummaryResult:
233    """Return a one-sentence TLDR and 2-4 key contributions."""
234    return _get_provider().summarize(content)
235
236
237def find_related(
238    content: PaperContent,
239    candidates: list[tuple[str, str]],   # [(paper_id, abstract), ...]
240    threshold: int = 5,
241) -> list[str]:
242    """
243    Return IDs of the most conceptually related papers from candidates.
244    Useful for adding semantic edges to the graph beyond shared category/author.
245    """
246    return _get_provider().find_related(content, candidates, threshold)