lookup.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187
  1. #!/usr/bin/env python3
  2. """
  3. Research Lookup Tool for Claude Code
  4. Performs research queries using Perplexity Sonar Pro Search via OpenRouter.
  5. """
  6. import os
  7. import sys
  8. import json
  9. from typing import Dict, List, Optional
  10. # Import the main research lookup class
  11. sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'scripts'))
  12. from research_lookup import ResearchLookup
  13. def format_response(result: Dict) -> str:
  14. """Format the research result for display."""
  15. if not result["success"]:
  16. return f"❌ Research lookup failed: {result['error']}"
  17. response = result["response"]
  18. citations = result["citations"]
  19. sources = result.get("sources", [])
  20. # Format the output for Claude Code
  21. output = f"""🔍 **Research Results**
  22. **Query:** {result['query']}
  23. **Model:** {result['model']}
  24. **Timestamp:** {result['timestamp']}
  25. **Note:** Results prioritized by citation count, venue prestige, and author reputation
  26. ---
  27. {response}
  28. """
  29. # Display API-provided sources with venue/citation info
  30. if sources:
  31. output += f"\n📚 **Sources ({len(sources)}):**\n"
  32. output += "_Prioritized by venue quality and citation impact_\n\n"
  33. for i, source in enumerate(sources, 1):
  34. title = source.get("title", "Untitled")
  35. url = source.get("url", "")
  36. date = source.get("date", "")
  37. snippet = source.get("snippet", "")
  38. # Format source entry with available metadata
  39. date_str = f" ({date})" if date else ""
  40. output += f"{i}. **{title}**{date_str}\n"
  41. # Add venue indicator if detectable from URL
  42. venue_indicator = _detect_venue_tier(url)
  43. if venue_indicator:
  44. output += f" 📊 Venue: {venue_indicator}\n"
  45. if url:
  46. output += f" 🔗 {url}\n"
  47. if snippet:
  48. output += f" _{snippet[:150]}{'...' if len(snippet) > 150 else ''}_\n"
  49. output += "\n"
  50. # Display extracted citations (DOIs, etc.)
  51. if citations:
  52. doi_citations = [c for c in citations if c.get("type") == "doi"]
  53. url_citations = [c for c in citations if c.get("type") == "url"]
  54. if doi_citations:
  55. output += f"\n🔗 **DOI References ({len(doi_citations)}):**\n"
  56. for i, citation in enumerate(doi_citations, 1):
  57. output += f"{i}. DOI: {citation.get('doi', '')} → {citation.get('url', '')}\n"
  58. if url_citations:
  59. output += f"\n🌐 **Additional URLs ({len(url_citations)}):**\n"
  60. for i, citation in enumerate(url_citations, 1):
  61. url = citation.get('url', '')
  62. venue = _detect_venue_tier(url)
  63. venue_str = f" [{venue}]" if venue else ""
  64. output += f"{i}. {url}{venue_str}\n"
  65. if result.get("usage"):
  66. usage = result["usage"]
  67. output += f"\n**Usage:** {usage.get('total_tokens', 'N/A')} tokens"
  68. return output
  69. def _detect_venue_tier(url: str) -> Optional[str]:
  70. """Detect venue tier from URL to indicate source quality."""
  71. if not url:
  72. return None
  73. url_lower = url.lower()
  74. # Tier 1 - Premier venues
  75. tier1_indicators = {
  76. "nature.com": "Nature (Tier 1)",
  77. "science.org": "Science (Tier 1)",
  78. "cell.com": "Cell Press (Tier 1)",
  79. "nejm.org": "NEJM (Tier 1)",
  80. "thelancet.com": "Lancet (Tier 1)",
  81. "jamanetwork.com": "JAMA (Tier 1)",
  82. "pnas.org": "PNAS (Tier 1)",
  83. }
  84. # Tier 2 - High-impact specialized
  85. tier2_indicators = {
  86. "neurips.cc": "NeurIPS (Tier 2 - Top ML)",
  87. "icml.cc": "ICML (Tier 2 - Top ML)",
  88. "openreview.net": "Top ML Conference (Tier 2)",
  89. "aacrjournals.org": "AACR Journals (Tier 2)",
  90. "ahajournals.org": "AHA Journals (Tier 2)",
  91. "bloodjournal.org": "Blood (Tier 2)",
  92. "jci.org": "JCI (Tier 2)",
  93. }
  94. # Tier 3 - Respected academic sources
  95. tier3_indicators = {
  96. "springer.com": "Springer",
  97. "wiley.com": "Wiley",
  98. "elsevier.com": "Elsevier",
  99. "oup.com": "Oxford University Press",
  100. "arxiv.org": "arXiv (Preprint)",
  101. "biorxiv.org": "bioRxiv (Preprint)",
  102. "medrxiv.org": "medRxiv (Preprint)",
  103. "pubmed": "PubMed",
  104. "ncbi.nlm.nih.gov": "NCBI/PubMed",
  105. "ieee.org": "IEEE",
  106. "acm.org": "ACM",
  107. }
  108. for domain, label in tier1_indicators.items():
  109. if domain in url_lower:
  110. return label
  111. for domain, label in tier2_indicators.items():
  112. if domain in url_lower:
  113. return label
  114. for domain, label in tier3_indicators.items():
  115. if domain in url_lower:
  116. return label
  117. return None
  118. def main():
  119. """Main entry point for Claude Code tool."""
  120. # Check for API key
  121. if not os.getenv("OPENROUTER_API_KEY"):
  122. print("❌ Error: OPENROUTER_API_KEY environment variable not set")
  123. print("Please set it in your .env file or export it:")
  124. print(" export OPENROUTER_API_KEY='your_openrouter_api_key'")
  125. return 1
  126. # Get query from command line arguments
  127. if len(sys.argv) < 2:
  128. print("❌ Error: No query provided")
  129. print("Usage: python lookup.py 'your research query here'")
  130. return 1
  131. query = " ".join(sys.argv[1:])
  132. try:
  133. # Initialize research tool
  134. research = ResearchLookup()
  135. # Perform lookup
  136. print(f"🔍 Researching: {query}")
  137. result = research.lookup(query)
  138. # Format and output result
  139. formatted_output = format_response(result)
  140. print(formatted_output)
  141. # Return success code
  142. return 0 if result["success"] else 1
  143. except Exception as e:
  144. print(f"❌ Error: {str(e)}")
  145. return 1
  146. if __name__ == "__main__":
  147. exit(main())