--- a
+++ b/api.py
@@ -0,0 +1,244 @@
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from aitrika.engine.online_aitrika import OnlineAItrika
+from aitrika.llm.openai import OpenAILLM
+from aitrika.utils.text_parser import generate_documents
+from dotenv import load_dotenv
+from fastapi.openapi.utils import get_openapi
+import json
+import os
+import uvicorn
+
+app = FastAPI(title="AItrika API", description="AItrika API", version="0.1.0")
+load_dotenv()
+
+
+class PubMedRequest(BaseModel):
+ pubmed_id: int
+
+
+class QueryRequest(BaseModel):
+ pubmed_id: int
+ query: str
+
+
+@app.post(
+ "/associations",
+ summary="Get associations from a PubMed article",
+ responses={
+ 200: {"description": "Associations retrieved successfully"},
+ 422: {"description": "Invalid PubMed ID provided"},
+ },
+)
+def get_associations(request: PubMedRequest):
+ """
+ Extracts associations from a PubMed article based on its ID.
+
+ Args:
+ request (PubMedRequest): An object containing the PubMed ID.
+
+ Returns:
+ dict: A dictionary containing the extracted associations.
+
+ Raises:
+ HTTPException: If the PubMed ID is invalid or the article is not found.
+ """
+ try:
+ engine = OnlineAItrika(pubmed_id=request.pubmed_id)
+ associations = engine.extract_associations()
+ return JSONResponse(content={"associations": associations}, status_code=200)
+ except ValueError as e:
+ raise HTTPException(status_code=422, detail=str(e))
+ except Exception:
+ raise HTTPException(status_code=500, detail="An unexpected error occurred")
+
+
+@app.post(
+ "/abstract",
+ summary="Get abstract of a PubMed article",
+ responses={
+ 200: {"description": "Abstract retrieved successfully"},
+ 422: {"description": "Invalid PubMed ID provided"},
+ },
+)
+def get_abstract(request: PubMedRequest):
+ """
+ Retrieves the abstract of a PubMed article based on its ID.
+
+ Args:
+ request (PubMedRequest): An object containing the PubMed ID.
+
+ Returns:
+ dict: A dictionary containing the abstract of the article.
+
+ Raises:
+ HTTPException: If the PubMed ID is invalid or the article is not found.
+ """
+ try:
+ engine = OnlineAItrika(pubmed_id=request.pubmed_id)
+ abstract = engine.extract_abstract()
+ return JSONResponse(content={"abstract": abstract}, status_code=200)
+ except ValueError as e:
+ raise HTTPException(status_code=422, detail=str(e))
+ except Exception:
+ raise HTTPException(status_code=500, detail="An unexpected error occurred")
+
+
+@app.post(
+ "/query",
+ summary="Query a PubMed article",
+ responses={
+ 200: {"description": "Query result retrieved successfully"},
+ 422: {"description": "Invalid PubMed ID or query provided"},
+ },
+)
+def query_document(request: QueryRequest):
+ """
+ Queries a PubMed article with a user-provided question.
+
+ Args:
+ request (QueryRequest): An object containing the PubMed ID and the query.
+
+ Returns:
+ dict: A dictionary containing the result of the query.
+
+ Raises:
+ HTTPException: If the PubMed ID is invalid, the article is not found, or the query fails.
+ """
+ try:
+ engine = OnlineAItrika(pubmed_id=request.pubmed_id)
+ abstract = engine.extract_abstract()
+ documents = generate_documents(content=abstract)
+ llm = OpenAILLM(
+ documents=documents,
+ api_key=os.getenv("OPENAI_API_KEY"),
+ )
+ result = llm.query(query=request.query)
+ return JSONResponse(content={"result": result}, status_code=200)
+ except ValueError as e:
+ raise HTTPException(status_code=422, detail=str(e))
+ except Exception:
+ raise HTTPException(status_code=500, detail="An unexpected error occurred")
+
+
+@app.post(
+ "/results",
+ summary="Get results from a PubMed article",
+ responses={
+ 200: {"description": "Results retrieved successfully"},
+ 422: {"description": "Invalid PubMed ID provided"},
+ },
+)
+def get_results(request: PubMedRequest):
+ """
+ Extracts results from a PubMed article based on its ID.
+
+ Args:
+ request (PubMedRequest): An object containing the PubMed ID.
+
+ Returns:
+ dict: A dictionary containing the extracted results.
+
+ Raises:
+ HTTPException: If the PubMed ID is invalid or the article is not found.
+ """
+ try:
+ engine = OnlineAItrika(pubmed_id=request.pubmed_id)
+ llm = OpenAILLM(
+ documents=generate_documents(content=engine.extract_abstract()),
+ api_key=os.getenv("OPENAI_API_KEY"),
+ )
+ results = engine.extract_results(llm=llm)
+ return JSONResponse(content={"results": results}, status_code=200)
+ except ValueError as e:
+ raise HTTPException(status_code=422, detail=str(e))
+ except Exception:
+ raise HTTPException(status_code=500, detail="An unexpected error occurred")
+
+
+@app.post(
+ "/participants",
+ summary="Get number of participants from a PubMed article",
+ responses={
+ 200: {"description": "Number of participants retrieved successfully"},
+ 422: {"description": "Invalid PubMed ID provided"},
+ },
+)
+def get_number_of_participants(request: PubMedRequest):
+ """
+ Extracts the number of participants from a PubMed article based on its ID.
+
+ Args:
+ request (PubMedRequest): An object containing the PubMed ID.
+
+ Returns:
+ dict: A dictionary containing the number of participants.
+
+ Raises:
+ HTTPException: If the PubMed ID is invalid or the article is not found.
+ """
+ try:
+ engine = OnlineAItrika(pubmed_id=request.pubmed_id)
+ llm = OpenAILLM(
+ documents=generate_documents(content=engine.extract_abstract()),
+ api_key=os.getenv("OPENAI_API_KEY"),
+ )
+ number_of_participants = engine.extract_number_of_participants(llm=llm)
+ return JSONResponse(
+ content={"number_of_participants": number_of_participants}, status_code=200
+ )
+ except ValueError as e:
+ raise HTTPException(status_code=422, detail=str(e))
+ except Exception:
+ raise HTTPException(status_code=500, detail="An unexpected error occurred")
+
+
+@app.post(
+ "/outcomes",
+ summary="Get outcomes from a PubMed article",
+ responses={
+ 200: {"description": "Outcomes retrieved successfully"},
+ 422: {"description": "Invalid PubMed ID provided"},
+ },
+)
+def get_outcomes(request: PubMedRequest):
+ """
+ Extracts outcomes from a PubMed article based on its ID.
+
+ Args:
+ request (PubMedRequest): An object containing the PubMed ID.
+
+ Returns:
+ dict: A dictionary containing the extracted outcomes.
+
+ Raises:
+ HTTPException: If the PubMed ID is invalid or the article is not found.
+ """
+ try:
+ engine = OnlineAItrika(pubmed_id=request.pubmed_id)
+ llm = OpenAILLM(
+ documents=generate_documents(content=engine.extract_abstract()),
+ api_key=os.getenv("OPENAI_API_KEY"),
+ )
+ outcomes = engine.extract_outcomes(llm=llm)
+ return JSONResponse(content={"outcomes": outcomes}, status_code=200)
+ except ValueError as e:
+ raise HTTPException(status_code=422, detail=str(e))
+ except Exception:
+ raise HTTPException(status_code=500, detail="An unexpected error occurred")
+
+
+if __name__ == "__main__":
+ with open("docs/api-reference/openapi.json", "w") as f:
+ json.dump(
+ get_openapi(
+ title=app.title,
+ version=app.version,
+ openapi_version=app.openapi_version,
+ description=app.description,
+ routes=app.routes,
+ ),
+ f,
+ )
+ uvicorn.run(app, host="0.0.0.0", port=8000)
