academia-mcp 1.1.4__py3-none-any.whl → 1.2.1__py3-none-any.whl

academia_mcp/llm.py

@@ -0,0 +1,38 @@
+import os
+from typing import List, Dict, Any
+
+from pydantic import BaseModel
+from openai import AsyncOpenAI
+from openai.types.chat.chat_completion_message import ChatCompletionMessage
+
+
+class ChatMessage(BaseModel):  # type: ignore
+    role: str
+    content: str | List[Dict[str, Any]]
+
+
+ChatMessages = List[ChatMessage]
+
+
+async def llm_acall(model_name: str, prompt: str) -> str:
+    key = os.getenv("OPENROUTER_API_KEY", "")
+    assert key, "Please set OPENROUTER_API_KEY in the environment variables"
+    base_url = os.getenv("BASE_URL", "https://openrouter.ai/api/v1")
+
+    messages: ChatMessages = [
+        ChatMessage(role="user", content=prompt),
+    ]
+    client = AsyncOpenAI(base_url=base_url, api_key=key)
+    response: ChatCompletionMessage = (
+        (
+            await client.chat.completions.create(
+                model=model_name,
+                messages=messages,
+                temperature=0.0,
+            )
+        )
+        .choices[0]
+        .message
+    )
+    assert response.content, "Response content is None"
+    return response.content

academia_mcp/server.py

@@ -15,7 +15,12 @@ from .tools.document_qa import document_qa
 from .tools.md_to_pdf import md_to_pdf
 from .tools.web_search import web_search, tavily_web_search, exa_web_search, brave_web_search
 from .tools.visit_webpage import visit_webpage
-from .tools.bitflip import propose_improvement_idea, extract_bitflip_info
+from .tools.bitflip import (
+    extract_bitflip_info,
+    generate_research_proposal,
+    score_research_proposals,
+)
+
 
 load_dotenv()
 
@@ -53,6 +58,9 @@ def run(
     server.add_tool(anthology_search)
     server.add_tool(md_to_pdf)
     server.add_tool(visit_webpage)
+    server.add_tool(extract_bitflip_info)
+    server.add_tool(generate_research_proposal)
+    server.add_tool(score_research_proposals)
 
     if os.getenv("TAVILY_API_KEY"):
         server.add_tool(tavily_web_search)
@@ -64,8 +72,6 @@ def run(
         server.add_tool(web_search)
     if os.getenv("OPENROUTER_API_KEY"):
         server.add_tool(document_qa)
-        server.add_tool(propose_improvement_idea)
-        server.add_tool(extract_bitflip_info)
 
     if port is None:
         port = find_free_port()

academia_mcp/tools/__init__.py

@@ -7,7 +7,7 @@ from .document_qa import document_qa
 from .md_to_pdf import md_to_pdf
 from .web_search import web_search, tavily_web_search, exa_web_search, brave_web_search
 from .visit_webpage import visit_webpage
-from .bitflip import propose_improvement_idea, extract_bitflip_info
+from .bitflip import extract_bitflip_info, generate_research_proposal, score_research_proposals
 
 
 __all__ = [
@@ -24,6 +24,7 @@ __all__ = [
     "exa_web_search",
     "brave_web_search",
     "visit_webpage",
-    "propose_improvement_idea",
     "extract_bitflip_info",
+    "generate_research_proposal",
+    "score_research_proposals",
 ]

academia_mcp/tools/bitflip.py

@@ -4,15 +4,24 @@
 import json
 import os
 import random
-from typing import List, Dict, Any
+from typing import List, Optional, Any
 
-from openai import OpenAI
 from pydantic import BaseModel
-from openai.types.chat import ChatCompletionMessage
 from datasets import load_dataset  # type: ignore
 
 from academia_mcp.tools.arxiv_download import arxiv_download
 from academia_mcp.utils import extract_json, encode_prompt
+from academia_mcp.llm import llm_acall
+
+
+class ProposalDataset:
+    dataset: Optional[List[Any]] = None
+
+    @classmethod
+    def get_dataset(cls) -> List[Any]:
+        if cls.dataset is None:
+            cls.dataset = list(load_dataset("UniverseTBD/hypogen-dr1")["train"])
+        return cls.dataset
 
 
 EXTRACT_PROMPT = """
@@ -85,24 +94,98 @@ Try to be as specific as possible.
 Now, please propose a chain of reasoning that leads to an improvement idea for this Bit:
 {{bit}}
 
+{% if additional_context %}Additional context:
+{{additional_context}}{% endif %}
+
+Finalize your idea by providing the idea details:
+- Abstract: An abstract that summarizes the proposal in conference format (approximately 250 words).
+- Experiments: A list of experiments that would be conducted to validate the proposal. Ensure these are simple and feasible. Be specific in exactly how you would test the hypothesis, and detail precise algorithmic changes. Include the evaluation metrics you would use.
+- Risks and limitations: A list of potential risks and limitations of the proposal.
+
 Return only the JSON object in this exact format (no extra text):
 {
     "chain_of_reasoning": "Chain of reasoning that leads to an improvement idea for this Bit. At least 5 sentences.",
     "flip": "Innovative approach or solution, in at least two sentences",
-    "spark": "4-6 word summary"
+    "spark": "4-6 word summary",
+    "abstract": "An abstract that summarizes the proposal in conference format (approximately 250 words).",
+    "experiments": ["...", "..."],
+    "risks_and_limitations": "A list of potential risks and limitations of the proposal."
 }
 """
 
 
-class ChatMessage(BaseModel):  # type: ignore
-    role: str
-    content: str | List[Dict[str, Any]]
+SCORE_PROMPT = """
+You are a highly advanced research assistant.
+You are given a list of research proposals.
+Your task is to score the proposals.
+
+Proposals:
+{% for proposal in proposals %}
+- Proposal ID: {{proposal["proposal_id"]}}
+- Spark: {{proposal["spark"]}}
+- Abstract: {{proposal["abstract"]}}
+- Experiments: {{proposal["experiments"]}}
+- Risks and limitations: {{proposal["risks_and_limitations"]}}
+{% endfor %}
+
+Here are the criteria:
+- "Strengths": A list of strengths of the proposal.
+- "Weaknesses": A list of weaknesses of the proposal.
+- "Novelty": Is the proposal novel? A rating from 1 to 4 (low, medium, high, very high).
+- "Clarity": Is the proposal clear? A rating from 1 to 4 (low, medium, high, very high).
+- "Significance": Is the proposal significant? A rating from 1 to 4 (low, medium, high, very high).
+- "Feasibility": Is the proposal feasible and easy to implement? A rating from 1 to 4 (low, medium, high, very high).
+- "Soundness": Is the proposal sound? A rating from 1 to 4 (poor, fair, good, excellent).
+- "Overall": A rating from 1 to 10 (very strong reject to award quality).
+
+Return only scores for all proposals in this exact format (no extra text):
+[
+    {
+        "proposal_id": 0,
+        "spark": "...",
+        "strengths": ["...", "..."],
+        "weaknesses": ["...", "..."],
+        "novelty": 2,
+        "clarity": 2,
+        "significance": 2,
+        "feasibility": 2,
+        "soundness": 2,
+        "overall": 5
+    },
+    ...
+]
+"""
+
+
+class BitFlipInfo(BaseModel):  # type: ignore
+    bit: str
+    flip: str
+    spark: str
+
 
+class Proposal(BaseModel):  # type: ignore
+    proposal_id: Optional[int] = None
+    flip: str
+    spark: str
+    abstract: str
+    experiments: List[str]
+    risks_and_limitations: List[str]
 
-ChatMessages = List[ChatMessage]
 
+class ProposalScores(BaseModel):  # type: ignore
+    proposal_id: int
+    spark: str
+    strengths: List[str]
+    weaknesses: List[str]
+    novelty: int
+    clarity: int
+    significance: int
+    feasibility: int
+    soundness: int
+    overall: int
 
-def extract_bitflip_info(arxiv_id: str) -> str:
+
+async def extract_bitflip_info(arxiv_id: str) -> str:
     """
     Extracts the Bit-Flip information from the arXiv paper.
 
@@ -121,73 +204,79 @@ def extract_bitflip_info(arxiv_id: str) -> str:
     Args:
         arxiv_id: The arXiv ID of the paper to extract the Bit-Flip information from.
     """
-    base_url = os.getenv("BASE_URL", "https://openrouter.ai/api/v1")
-    key = os.getenv("OPENROUTER_API_KEY", "")
-    assert key, "Please set OPENROUTER_API_KEY in the environment variables"
     model_name = os.getenv("BITFLIP_MODEL_NAME", "deepseek/deepseek-chat-v3-0324")
-
     paper = arxiv_download(arxiv_id)
     abstract = json.loads(paper)["abstract"]
     prompt = encode_prompt(EXTRACT_PROMPT, abstract=abstract)
-    messages: ChatMessages = [
-        ChatMessage(role="user", content=prompt),
-    ]
-    client = OpenAI(base_url=base_url, api_key=key)
-    response: ChatCompletionMessage = (
-        client.chat.completions.create(
-            model=model_name,
-            messages=messages,
-            temperature=0.0,
-        )
-        .choices[0]
-        .message
-    )
-    assert response.content, "Response content is None"
-    result = extract_json(response.content)
-    return json.dumps(result, ensure_ascii=False)
+    content = await llm_acall(model_name=model_name, prompt=prompt)
+    result = extract_json(content)
+    bitflip_info: BitFlipInfo = BitFlipInfo.model_validate(result)
+    return str(bitflip_info.model_dump_json())
 
 
-def propose_improvement_idea(arxiv_id: str) -> str:
+async def generate_research_proposal(bit: str, additional_context: str = "") -> str:
     """
-    Proposes an improvement idea for the arXiv paper.
+    Proposes an improvement idea for the Bit.
 
-    Returns a JSON object in this format:
+    Args:
+        bit: The Bit to propose an improvement idea for. The bit is a technical limitation or conventional approach of some paper.
+        additional_context: Additional context to use when proposing the improvement idea.
+
+    Returns a JSON string with a research proposal in this format:
     {
-        "chain_of_reasoning": "Chain of reasoning that leads to an improvement idea.",
-        "flip": "Innovative approach or solution",
-        "spark": "4-6 word summary"
+        "proposal_id": ...,
+        "flip": "Innovative approach or solution, in at least two sentences",
+        "spark": "4-6 word summary",
+        "abstract": "An abstract that summarizes the proposal in conference format (approximately 250 words).",
+        "experiments": ["...", "..."],
+        "risks_and_limitations": "A list of potential risks and limitations of the proposal."
     }
     Use `json.loads` to deserialize the result if you want to get specific fields.
-
-    Args:
-        arxiv_id: The arXiv ID of the paper to propose an improvement idea for.
     """
-    base_url = os.getenv("BASE_URL", "https://openrouter.ai/api/v1")
-    key = os.getenv("OPENROUTER_API_KEY", "")
-    assert key, "Please set OPENROUTER_API_KEY in the environment variables"
     model_name = os.getenv("BITFLIP_MODEL_NAME", "deepseek/deepseek-chat-v3-0324")
+    examples = ProposalDataset.get_dataset()[:]
+    examples = random.choices(examples, k=4)
 
-    bitflip_info = json.loads(extract_bitflip_info(arxiv_id))
-    bit = bitflip_info["bit"]
+    prompt = encode_prompt(
+        IMPROVEMENT_PROMPT, bit=bit, examples=examples, additional_context=additional_context
+    )
+    content = await llm_acall(model_name=model_name, prompt=prompt)
+    result = extract_json(content)
+    proposal: Proposal = Proposal.model_validate(result)
+    proposal.proposal_id = random.randint(0, 1000000)
+    return str(proposal.model_dump_json())
 
-    examples = list(load_dataset("UniverseTBD/hypogen-dr1")["train"])
-    random.shuffle(examples)
-    examples = examples[:4]
 
-    prompt = encode_prompt(IMPROVEMENT_PROMPT, bit=bit, examples=examples)
-    messages: ChatMessages = [
-        ChatMessage(role="user", content=prompt),
+async def score_research_proposals(proposals: List[str]) -> str:
+    """
+    Scores a list of research proposals.
+    Use proposals obtained with the `generate_research_proposal` tool.
+
+    Returns a JSON string with a list of scores in this format:
+    [
+        {
+            "proposal_id": 0,
+            "spark": "...",
+            "strengths": ["...", "..."],
+            "weaknesses": ["...", "..."],
+            "novelty": 2,
+            "clarity": 2,
+            "significance": 2,
+            "feasibility": 2,
+            "soundness": 2,
+            "overall": 5
+        },
+        ...
     ]
-    client = OpenAI(base_url=base_url, api_key=key)
-    response: ChatCompletionMessage = (
-        client.chat.completions.create(
-            model=model_name,
-            messages=messages,
-            temperature=0.0,
-        )
-        .choices[0]
-        .message
-    )
-    assert response.content, "Response content is None"
-    result = extract_json(response.content)
-    return json.dumps(result, ensure_ascii=False)
+    Use `json.loads` to deserialize the result if you want to get specific fields.
+
+    Args:
+        proposals: A list of JSON strings with research proposals.
+    """
+    model_name = os.getenv("BITFLIP_MODEL_NAME", "deepseek/deepseek-chat-v3-0324")
+    proposals = [Proposal.model_validate_json(proposal) for proposal in proposals]
+    prompt = encode_prompt(SCORE_PROMPT, proposals=proposals)
+    content = await llm_acall(model_name=model_name, prompt=prompt)
+    scores = extract_json(content)
+    final_scores = [ProposalScores.model_validate(score) for score in scores]
+    return json.dumps([s.model_dump() for s in final_scores], ensure_ascii=False)

academia_mcp/tools/document_qa.py

@@ -1,18 +1,16 @@
 import os
-from typing import List, Any, Dict, cast
+import json
+from typing import List, Any, Dict
 from dotenv import load_dotenv
 
 from pydantic import BaseModel
-from openai import OpenAI
-from openai.types.chat import ChatCompletionMessageParam, ChatCompletionMessage
 
+from academia_mcp.llm import llm_acall
 
 load_dotenv()
 
-SYSTEM_PROMPT = (
-    "You are a helpful assistant that answers questions about documents accurately and concisely."
-)
-PROMPT = """Please answer the following questions based solely on the provided document.
+PROMPT = """You are a helpful assistant that answers questions about documents accurately and concisely.
+Please answer the following questions based solely on the provided document.
 If there is no answer in the document, output "There is no answer in the provided document".
 First cite ALL relevant document fragments, then provide a final answer.
 Answer all given questions one by one.
@@ -40,8 +38,8 @@ class ChatMessage(BaseModel):  # type: ignore
 ChatMessages = List[ChatMessage]
 
 
-def document_qa(
-    document: str,
+async def document_qa(
+    document: str | Dict[str, Any],
     question: str,
 ) -> str:
     """
@@ -58,39 +56,16 @@ def document_qa(
     Returns an answer to all questions based on the document content.
 
     Args:
-    question: Question (or questions) to be answered about the document.
-    document: The full text of the document to analyze.
+        question: Question (or questions) to be answered about the document.
+        document: The full text of the document to analyze.
     """
     assert question and question.strip(), "Please provide non-empty 'question'"
+    if isinstance(document, dict):
+        document = json.dumps(document)
+
     assert document and document.strip(), "Please provide non-empty 'document'"
 
-    base_url = os.getenv("BASE_URL", "https://openrouter.ai/api/v1")
-    key = os.getenv("OPENROUTER_API_KEY", "")
-    assert key, "Please set OPENROUTER_API_KEY in the environment variables"
     model_name = os.getenv("DOCUMENT_QA_MODEL_NAME", "deepseek/deepseek-chat-v3-0324")
-
-    messages: ChatMessages = [
-        ChatMessage(role="system", content=SYSTEM_PROMPT),
-        ChatMessage(
-            role="user",
-            content=PROMPT.format(question=question, document=document),
-        ),
-    ]
-
-    sdk_messages = [
-        cast(ChatCompletionMessageParam, m.model_dump(exclude_none=True)) for m in messages
-    ]
-    client = OpenAI(base_url=base_url, api_key=key)
-    response: ChatCompletionMessage = (
-        client.chat.completions.create(
-            model=model_name,
-            messages=sdk_messages,
-            temperature=0.0,
-        )
-        .choices[0]
-        .message
-    )
-
-    if response.content is None:
-        raise Exception("Response content is None")
-    return response.content.strip()
+    prompt = PROMPT.format(question=question, document=document)
+    content = await llm_acall(model_name=model_name, prompt=prompt)
+    return content.strip()

{academia_mcp-1.1.4.dist-info → academia_mcp-1.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: academia-mcp
-Version: 1.1.4
+Version: 1.2.1
 Summary: MCP server that provides different tools to search for scientific publications
 Author-email: Ilya Gusev <phoenixilya@gmail.com>
 Project-URL: Homepage, https://github.com/IlyaGusev/academia_mcp
@@ -31,6 +31,7 @@ Requires-Dist: pytest>=8.4.1
 Requires-Dist: openai>=1.97.1
 Requires-Dist: jinja2>=3.1.6
 Requires-Dist: datasets>=4.0.0
+Requires-Dist: pytest-asyncio>=1.1.0
 Dynamic: license-file
 
 # Academia MCP

{academia_mcp-1.1.4.dist-info → academia_mcp-1.2.1.dist-info}/RECORD

@@ -1,24 +1,25 @@
 academia_mcp/__init__.py,sha256=2Ru2I5u4cE7DrkkAsibDUEF1K6sYtqppb9VyFrRoQKI,94
 academia_mcp/__main__.py,sha256=rcmsOtJd3SA82exjrcGBuxuptcoxF8AXI7jNjiVq2BY,59
 academia_mcp/files.py,sha256=hI5dj4h0fX8V3DXKI_C8vs1fte2uc9gsBXC6prLV4o4,745
+academia_mcp/llm.py,sha256=o84FQNSbjjVSk9DlvFXWsUDiz5IOaavYU6kOqnPEG7E,1071
 academia_mcp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-academia_mcp/server.py,sha256=oiwai6-i_DZ8cCaOJuFRBmKyeLPj5vBTQlHxnmfAnYc,2406
+academia_mcp/server.py,sha256=q9bhacVm-8uuIMc_xSeymmVaIB8pQCqfTQx5GU8hhpM,2492
 academia_mcp/utils.py,sha256=tkHBho-NfzAR8rplFaiRYq4sBmQ9V3JPjDJTDrz58Xs,4041
-academia_mcp/tools/__init__.py,sha256=WUhrF0QHq6F3ljsK5Kc7jh1d0beVFeYVhRfKAjXK79M,872
+academia_mcp/tools/__init__.py,sha256=8_8QWPRYmPiUjdiNrQilEEnCRR-UBU7g-56jT52V3VQ,934
 academia_mcp/tools/anthology_search.py,sha256=_5s8EzdV7NQD_F3bjVH4XlKKHOJlFtWlQVrPbODuc3I,7847
 academia_mcp/tools/arxiv_download.py,sha256=xanzt77TZBQRngzGbKCRz4Hp-Mwfe_q-46eRW23TpVs,11219
 academia_mcp/tools/arxiv_search.py,sha256=pzM18qrF3QL03A53w003kE7hQi3s3QKtjgw0m7K88UY,8355
-academia_mcp/tools/bitflip.py,sha256=m6qWbPieXUmi-wSiIbBl1wsvEPwEL6m0f0XmM2xIIY0,7353
-academia_mcp/tools/document_qa.py,sha256=NrosUBgz__s8jCWpl4x17A4hMIPczFPFlfmR_OspKCQ,3005
+academia_mcp/tools/bitflip.py,sha256=u0hSOPWbnCDu2EbA_RkueX496SvTKz9QhZcXugshSfI,10949
+academia_mcp/tools/document_qa.py,sha256=lWtzRNFKwQpQO8vPNvYYbNnGh3LsBmGl6f9vSRIuqaw,2245
 academia_mcp/tools/hf_datasets_search.py,sha256=KiBkqT4rXjEN4oc1AWZOPnqN_Go90TQogY5-DUm3LQo,2854
 academia_mcp/tools/md_to_pdf.py,sha256=Ovc_-8j7gIZNEM1d0ZDH-8qbtfZLSaNmCm5DQjrtM0k,12810
 academia_mcp/tools/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 academia_mcp/tools/s2_citations.py,sha256=dqrBp76RrX1zH2XzcMAoWBbvbtyhxLeF-xnqOKD_JiM,4852
 academia_mcp/tools/visit_webpage.py,sha256=0zAZYeQxPDu0OjgAAvbMLZh0ttaS5q-_4WhgsEPrbsI,1542
 academia_mcp/tools/web_search.py,sha256=NAkbXdD9mKxsIXhWN32dRd_EiaB3G6ENy-n-bc7HAaQ,5448
-academia_mcp-1.1.4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-academia_mcp-1.1.4.dist-info/METADATA,sha256=g1qVsuRWZCAJ_EIjEdbG89Adq2FJ7Dy1infEIV8i8yQ,1862
-academia_mcp-1.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-academia_mcp-1.1.4.dist-info/entry_points.txt,sha256=gxkiKJ74w2FwJpSECpjA3XtCfI5ZfrM6N8cqnwsq4yY,51
-academia_mcp-1.1.4.dist-info/top_level.txt,sha256=CzGpRFsRRJRqWEb1e3SUlcfGqRzOxevZGaJWrtGF8W0,13
-academia_mcp-1.1.4.dist-info/RECORD,,
+academia_mcp-1.2.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+academia_mcp-1.2.1.dist-info/METADATA,sha256=mLcD1By5dg9L8hNMlybFg5IUjS4S7uzXggGshnTESV8,1899
+academia_mcp-1.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+academia_mcp-1.2.1.dist-info/entry_points.txt,sha256=gxkiKJ74w2FwJpSECpjA3XtCfI5ZfrM6N8cqnwsq4yY,51
+academia_mcp-1.2.1.dist-info/top_level.txt,sha256=CzGpRFsRRJRqWEb1e3SUlcfGqRzOxevZGaJWrtGF8W0,13
+academia_mcp-1.2.1.dist-info/RECORD,,