Use headless search endpoint for Paper Finder (#14)

Author: rodneykinney

.claude-plugin/marketplace.json

@@ -11,7 +11,7 @@
       "name": "asta",
       "source": "./",
       "description": "Paper search, citations, literature reports, and Semantic Scholar API tools",
-      "version": "0.4.0",
+      "version": "0.5.0",
       "author": {
         "name": "AI2 Asta Team"
       },

.claude-plugin/plugin.json

@@ -1,6 +1,6 @@
 {
   "name": "asta",
-  "version": "0.4.0",
+  "version": "0.5.0",
   "description": "Asta science tools for Claude Code - paper search, citations, and more",
   "author": {
     "name": "AI2 Asta Team"

DEVELOPER.md

@@ -220,14 +220,17 @@ from asta.core import AstaPaperFinder
 
 client = AstaPaperFinder()
 
-# Simple blocking search
+# Simple synchronous search using headless endpoint
 result = client.find_papers("query", timeout=300)
-# Returns: {widget_id, file_path, paper_count}
-
-# Non-blocking start
-thread_id = client.start_search("query")
-widget_id = client.get_widget_id(thread_id)
-results = client.poll_for_results(widget_id, timeout=300)
+# Returns: {query, widget, status, timestamp, paper_count}
+
+# With operation mode control
+result = client.find_papers(
+    "query",
+    timeout=300,
+    operation_mode="fast",  # "infer", "fast", or "diligent"
+    include_full_metadata=True
+)
 ```
 
 ### SemanticScholarClient

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "asta"
-version = "0.4.0"
+version = "0.5.0"
 description = "Asta CLI for scientific literature review"
 readme = "README.md"
 requires-python = ">=3.11"

src/asta/__init__.py

@@ -1,3 +1,3 @@
 """Asta - Science literature research tools"""
 
-__version__ = "0.4.0"
+__version__ = "0.5.0"

src/asta/core/client.py

@@ -7,149 +7,86 @@
 import time
 import urllib.error
 import urllib.request
-import uuid
 from pathlib import Path
 from typing import Any
 
 
 class AstaPaperFinder:
-    """Client for Asta Paper Finder API"""
+    """Client for Asta Paper Finder API using headless endpoint"""
 
-    def __init__(self, base_url: str = "REDACTED_ASTA_PROD_URL"):
+    def __init__(self, base_url: str = "REDACTED_MABOOL_WORKERS_URL"):
         self.base_url = base_url
-        self.mabool_url = "REDACTED_MABOOL_DEMO_URL"
-        self.user_id = str(uuid.uuid4())
         self.headers = {
-            "X-Anonymous-User-ID": self.user_id,
             "Content-Type": "application/json",
         }
 
     def _request(
         self, url: str, method: str = "GET", data: dict | None = None
-    ) -> dict[str, Any] | list:
+    ) -> dict[str, Any]:
         """Make an HTTP request and return JSON response"""
         body = json.dumps(data).encode() if data else None
         req = urllib.request.Request(
             url, data=body, headers=self.headers, method=method
         )
-        response = urllib.request.urlopen(req)
-        return json.loads(response.read())
-
-    def create_thread(self) -> str:
-        """Create a new thread"""
-        result = self._request(f"{self.base_url}/api/chat/thread", method="PUT")
-        return result["thread"]["key"]
-
-    def send_message(
-        self, text: str, thread_id: str, profile: str = "paper-finder-only"
-    ) -> dict[str, Any]:
-        """Send a message to the thread"""
-        return self._request(
-            f"{self.base_url}/api/chat/message",
-            method="POST",
-            data={"text": text, "thread_id": thread_id, "profile": profile},
-        )
-
-    def get_widget_id(self, thread_id: str, max_retries: int = 20) -> str | None:
-        """Get the widget ID from thread events"""
-        url = f"{self.base_url}/api/rest/thread/{thread_id}/event/widget_paper_finder"
-        for _ in range(max_retries):
-            try:
-                req = urllib.request.Request(url, headers=self.headers)
-                response = urllib.request.urlopen(req)
-                data = json.loads(response.read())
-                last_event = data.get("last_event")
-                if last_event and isinstance(last_event, dict):
-                    event_data = last_event.get("data")
-                    if event_data and isinstance(event_data, dict):
-                        widget_id = event_data.get("id")
-                        if widget_id:
-                            return widget_id
-            except urllib.error.HTTPError:
-                pass
-            time.sleep(2)
-        return None
-
-    def get_widget_results(self, widget_id: str) -> dict[str, Any] | list:
-        """Get widget results from mabool service"""
-        url = f"{self.mabool_url}/api/2/rounds/{widget_id}/result/widget"
-        req = urllib.request.Request(url, headers=self.headers)
-        response = urllib.request.urlopen(req)
-        return json.loads(response.read())
-
-    def poll_for_results(self, widget_id: str, timeout: int = 300):
-        """Poll for results until completion or timeout"""
-        start_time = time.time()
-        poll_interval = 2
-
-        while time.time() - start_time < timeout:
+        try:
+            response = urllib.request.urlopen(req)
+            return json.loads(response.read())
+        except urllib.error.HTTPError as e:
+            error_body = e.read().decode("utf-8")
             try:
-                result = self.get_widget_results(widget_id)
-
-                # Handle if result is a list - got the papers directly
-                if isinstance(result, list):
-                    return {
-                        "roundStatus": {"kind": "completed"},
-                        "results": result,
-                        "thread_id": None,
-                        "widget_id": widget_id,
-                    }
-
-                # Handle dict response with roundStatus
-                status = result.get("roundStatus", {}).get("kind", "unknown")
-
-                if status == "completed":
-                    return result
-                elif status == "failed":
-                    error = result.get("roundStatus", {}).get("error", "Unknown error")
-                    raise Exception(f"Paper finder failed: {error}")
-
-            except urllib.error.HTTPError as e:
-                if e.code != 404:
-                    raise
-
-            time.sleep(poll_interval)
-
-        raise TimeoutError(f"Timeout after {timeout} seconds")
-
-    def start_search(self, query: str) -> str:
-        """Start a paper search and return thread_id immediately (non-blocking)"""
-        thread_id = self.create_thread()
-        self.send_message(query, thread_id)
-        return thread_id
+                error_data = json.loads(error_body)
+                error_msg = error_data.get("detail", str(e))
+            except json.JSONDecodeError:
+                error_msg = error_body or str(e)
+            raise Exception(f"API request failed: {error_msg}") from e
 
     def find_papers(
-        self, query: str, timeout: int = 300, save_to_file: Path | None = None
+        self,
+        query: str,
+        timeout: int = 300,
+        save_to_file: Path | None = None,
+        operation_mode: str = "infer",
+        include_full_metadata: bool = True,
     ) -> dict[str, Any]:
-        """Complete workflow to find papers (blocking).
+        """Execute a one-shot paper search using the headless endpoint.
 
         Args:
             query: Search query
-            timeout: Maximum time to wait for results
+            timeout: Maximum time to wait for results (seconds)
             save_to_file: Optional path to save results. If None, no file is saved.
+            operation_mode: Search strategy - 'infer', 'fast', or 'diligent' (default: 'infer')
+            include_full_metadata: Whether to return full paper details (default: True)
 
         Returns:
-            Complete search results including widget data
+            Complete search results with papers
         """
-        thread_id = self.start_search(query)
+        url = f"{self.base_url}/api/3/headless/paper-search"
+
+        request_body = {
+            "query": query,
+            "operation_mode": operation_mode,
+            "include_full_metadata": include_full_metadata,
+            "timeout_seconds": timeout,
+        }
 
-        # Get widget ID
-        widget_id = self.get_widget_id(thread_id)
-        if not widget_id:
-            raise Exception("Failed to get widget ID after retries")
+        # Make the synchronous request
+        result = self._request(url, method="POST", data=request_body)
 
-        # Poll for results
-        widget_result = self.poll_for_results(widget_id, timeout)
+        # Check for errors
+        if "error" in result and result["error"]:
+            error = result["error"]
+            raise Exception(f"Paper search failed: {error}")
 
-        papers = widget_result.get("results", [])
+        papers = result.get("papers", [])
 
-        # Build complete search data
+        # Build search data in format compatible with existing models
         search_data = {
             "query": query,
-            "thread_id": thread_id,
-            "widget_id": widget_id,
+            "widget": {
+                "results": papers,
+                "response_text": result.get("response_text", ""),
+            },
             "status": "completed",
-            "widget": widget_result,
             "timestamp": time.time(),
             "paper_count": len(papers),
         }

src/asta/literature/find.py

@@ -19,7 +19,13 @@
     default=300,
     help="Maximum time to wait for results (seconds)",
 )
-def find(query: str, timeout: int):
+@click.option(
+    "--mode",
+    type=click.Choice(["infer", "fast", "diligent"]),
+    default="infer",
+    help="Search strategy: infer (auto-detect), fast (quick results), or diligent (comprehensive)",
+)
+def find(query: str, timeout: int, mode: str):
     """Find papers matching QUERY using Asta Paper Finder.
 
     Saves results to .asta/literature/find/ with an auto-generated filename.
@@ -31,10 +37,18 @@ def find(query: str, timeout: int):
 
         # With custom timeout
         asta literature find "transformers" --timeout 60
+
+        # Use fast mode for quick results
+        asta literature find "deep learning" --mode fast
+
+        # Use diligent mode for comprehensive search
+        asta literature find "neural networks" --mode diligent
     """
     try:
         client = AstaPaperFinder()
-        raw_result = client.find_papers(query, timeout=timeout, save_to_file=None)
+        raw_result = client.find_papers(
+            query, timeout=timeout, save_to_file=None, operation_mode=mode
+        )
 
         # Transform to literature search result format
         literature_result = LiteratureSearchResult(

src/asta/literature/models.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 
 class Author(BaseModel):
@@ -62,27 +62,53 @@ class CitationContext(BaseModel):
 class Paper(BaseModel):
     """Paper search result with relevance judgements"""
 
-    corpusId: int
+    model_config = ConfigDict(populate_by_name=True)
+
+    # Use validation_alias to accept snake_case from API
+    corpusId: int = Field(validation_alias="corpus_id")
     title: str
     abstract: str | None = None
     year: int | None = None
     authors: list[Author] = Field(default_factory=list)
     venue: str | None = None
     journal: dict[str, Any] | None = None
     url: str | None = None
-    publicationDate: str | None = None
-    citationCount: int | None = None
+    publicationDate: str | None = Field(
+        default=None, validation_alias="publication_date"
+    )
+    citationCount: int | None = Field(default=None, validation_alias="citation_count")
     categories: list[str] = Field(default_factory=list)
 
     # Asta Paper Finder specific fields
-    relevanceScore: float
-    relevanceJudgement: RelevanceJudgement | None = None
+    relevanceScore: float = Field(validation_alias="relevance_score")
+    relevanceJudgement: RelevanceJudgement | None = Field(
+        default=None, validation_alias="relevance_judgement"
+    )
     snippets: list[Snippet] = Field(default_factory=list)
-    citationContexts: list[CitationContext] = Field(default_factory=list)
+    citationContexts: list[CitationContext] = Field(
+        default_factory=list, validation_alias="citation_contexts"
+    )
 
     # Legal/filtering fields
-    legalToShow: bool = True
-    numOfOmittedCitationContextsDueLegal: int = 0
+    legalToShow: bool = Field(default=True, validation_alias="legal_to_show")
+    numOfOmittedCitationContextsDueLegal: int = Field(
+        default=0, validation_alias="num_of_omitted_citation_contexts_due_legal"
+    )
+
+    @field_validator("authors", mode="before")
+    @classmethod
+    def convert_author_strings(cls, v):
+        """Convert author strings to Author objects if needed."""
+        if not isinstance(v, list):
+            return v
+        result = []
+        for author in v:
+            if isinstance(author, str):
+                # Convert string to Author dict
+                result.append({"name": author, "id": ""})
+            else:
+                result.append(author)
+        return result
 
 
 class LiteratureSearchResult(BaseModel):

tests/test_cli.py

@@ -140,7 +140,40 @@ def test_find_custom_timeout(self, runner):
 
         assert result.exit_code == 0
         mock_instance.find_papers.assert_called_once_with(
-            "test query", timeout=60, save_to_file=None
+            "test query", timeout=60, save_to_file=None, operation_mode="infer"
+        )
+
+    def test_find_with_mode_option(self, runner):
+        """Test find command with different operation modes."""
+        mock_result = {
+            "query": "test query",
+            "status": "completed",
+            "paper_count": 1,
+            "widget": {
+                "results": [
+                    {
+                        "corpusId": 123,
+                        "title": "Test Paper",
+                        "relevanceScore": 0.9,
+                        "authors": [],
+                    }
+                ]
+            },
+        }
+
+        with patch("asta.literature.find.AstaPaperFinder") as MockFinder:
+            mock_instance = MagicMock()
+            mock_instance.find_papers.return_value = mock_result
+            MockFinder.return_value = mock_instance
+
+            # Test fast mode
+            result = runner.invoke(
+                cli, ["literature", "find", "test query", "--mode", "fast"]
+            )
+
+        assert result.exit_code == 0
+        mock_instance.find_papers.assert_called_with(
+            "test query", timeout=300, save_to_file=None, operation_mode="fast"
         )