add calling openai assistant

requirements.txt

@@ -1,3 +1,4 @@
 fastapi
 uvicorn[standard]
 dashscope
+openai

src/talk_to_your_manual/__init__.py

@@ -1,9 +1,10 @@
 from fastapi import APIRouter
-from . import use_aliyun
+from . import use_aliyun, use_openai
 
 
 router = APIRouter(
     prefix="/talk-to-your-manual",
     tags=["Talk To Your Manual"],
 )
-router.include_router(use_aliyun.router)
+router.include_router(use_aliyun.router)
+router.include_router(use_openai.router)

src/talk_to_your_manual/use_aliyun.py

@@ -5,8 +5,8 @@ from fastapi.responses import JSONResponse
 import os
 
 router = APIRouter()
-API_KEY = os.environ.get("API_KEY")
-APP_ID = os.environ.get("APP_ID")
+API_KEY = os.environ.get("aliyun_api_key")
+APP_ID = os.environ.get("aliyun_app_id")
 
 session_id = None
 
@@ -49,6 +49,7 @@ async def call_aliyun(prompt: str):
         prompt=prompt,
         session_id=session_id,
     )
+    print(response.output.session_id)
     if response.status_code == HTTPStatus.OK:
         last_session_id = session_id
         session_id = response.output.session_id

src/talk_to_your_manual/use_openai.py

@@ -0,0 +1,70 @@
+from fastapi.responses import JSONResponse
+from fastapi import APIRouter, status
+from openai import OpenAI
+import os
+
+API_KEY = os.environ.get("openai_api_key")
+ASSISTANT_ID = os.environ.get("openai_assistant_id")
+router = APIRouter()
+client = OpenAI(api_key=API_KEY)
+
+
+@router.get("/openai")
+async def call_openai(prompt: str):
+    thread = client.beta.threads.create(
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": prompt,
+                    }
+                ],
+            }
+        ]
+    )
+    run = client.beta.threads.runs.create_and_poll(
+        thread_id=thread.id,
+        assistant_id=ASSISTANT_ID,
+    )
+    start_time = run.started_at
+    while True:
+        messages = list(
+            client.beta.threads.messages.list(thread_id=thread.id, run_id=run.id)
+        )
+        if messages:
+            print("received {} messages".format(len(messages)))
+            break
+        end_time = run.completed_at
+        if end_time - start_time > 60:
+            break
+    if not messages:
+        return JSONResponse(
+            status_code=status.HTTP_408_REQUEST_TIMEOUT,
+            content={
+                "code": run.code,
+                "message": run.message,
+            },
+        )
+    message_content = messages[0].content[0].text
+    annotations = message_content.annotations
+    citations = []
+    for index, annotation in enumerate(annotations):
+        message_content.value = message_content.value.replace(
+            annotation.text, f"[{index}]"
+        )
+        file_citation = getattr(annotation, "file_citation", None)
+        if file_citation:
+            cited_file = client.files.retrieve(file_citation.file_id)
+            citations.append(f"[{index}] {cited_file.filename}")
+    return JSONResponse(
+        status_code=status.HTTP_200_OK,
+        content={
+            "assistant_id": ASSISTANT_ID,
+            "thread_id": thread.id,
+            "run_id": run.id,
+            "answer": message_content.value,
+            "citations": citations,
+        },
+    )