mvp ok

app.py

@@ -1,83 +1,29 @@
 # uvicorn app:app --host 0.0.0.0 --port 7860 --reload
 
+from fastapi import FastAPI, HTTPException, Response, Depends
+from fastapi.responses import HTMLResponse
 from pathlib import Path
+from dotenv import load_dotenv
+import os
+
 from global_state import set
 
+load_dotenv()
+
 # 获取当前文件的父目录的绝对路径，即：project_root
 parent_dir = Path(__file__).resolve().parent
 set('project_root', parent_dir)
 
-
-from fastapi import FastAPI, HTTPException, Response, Depends
-from fastapi.responses import HTMLResponse
-# from starlette.requests import Request
-
 # from routers.webtools_v1 import router as webtools_router
 from routers.users_v1 import router as users_router
-# from routers.openai_v1 import router as openai_router
-
-# from db import Db
-
-from dotenv import load_dotenv
-import os
-
-# 当在本地使用.env 文件时加载.env文件
-load_dotenv()
-
-# 获取环境变量，本地使用 load_dotenv 和 hf 里直接使用环境变量配置时，都可使用下面的语句 
-DB_PATH = os.getenv("DB_PATH")
-
-# 依赖注入，获取 DB_PATH
-def get_db_path():
-    return DB_PATH
+from routers.openai_v1 import router as openai_router
 
 app = FastAPI()
 
 # app.include_router(webtools_router, prefix="/airs/v1", tags=["webtools"])
 app.include_router(users_router, prefix="/airs/v1", tags=["users"])
-# app.include_router(openai_router, prefix="/airs/v1", tags=["openai"])
+app.include_router(openai_router, prefix="/airs/v1", tags=["openai"])
 
 @app.get("/")
-def greet_json(db_path: str = Depends(get_db_path)):
+def greet_json():
     return {"Hello": "World!"}  
-# @app.get("/init_db")
-# def init_db():
-#     sql = """
-# CREATE TABLE api_keys (
-#     id INTEGER PRIMARY KEY AUTOINCREMENT,
-#     api_key TEXT NOT NULL,
-#     type TEXT NOT NULL,
-#     status INTEGER NOT NULL,
-#     idx INTEGER NOT NULL,
-#     dest_api_key TEXT NOT NULL
-# );
-# """
-    # Db('api_keys.db').execute_query(sql)
-
-# def create_user_api_key(user_id, api_key, type, status, idx, dest_api_key):
-#     sql = f"""
-# INSERT INTO api_keys (api_key, type, status, idx, dest_api_key)
-# """
-#     Db('api_keys.db').execute_query(sql)
-
-# create_user_api_key('i_am_tanbushi', 'llm', '1', '0', 'dest_api_key')
-
-# @app.get("/list_files")
-# def list_files():
-#     # return "cur_path"
-#     directory = os.getcwd()
-
-#     # 遍历目录树
-#     retstr="第一行\n第二行\n第三行\n"
-#     # retstr = retstr + '1' + "\n"
-#     # retstr = retstr + '2' + "\n"
-#     # retstr = retstr + '3' + "\n"
-#     print(retstr)
-#     # for dirpath, dirnames, filenames in os.walk(directory):
-#     #     for filename in filenames:
-#     #         retstr = retstr + os.path.join(dirpath, filename) +'\\n'
-#     #         print(os.path.join(dirpath, filename))
-
-#     return HTMLResponse(f"<pre>{retstr}</pre>")
-#     return HTMLResponse(content=f"<pre>{retstr}</pre>", media_type="text/html")
-#     return retstr

auth.py

@@ -8,11 +8,9 @@ from db_model.user import UserModel
 security = HTTPBearer()
 
 def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security)):
-    token = credentials.credentials
     token = credentials.credentials
     # 假设你有一个函数来验证Token并返回用户
     user = validate_token(token)
-    print(f"\n\n\n\n{user}")
     if user is None:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
@@ -26,7 +24,8 @@ def validate_token(token: str):
     db_module_filename = f"{get('project_root')}/db/cloudflare.py"
     query = f"SELECT * FROM users where api_key='{token}'"
     response = TbsDb(db_module_filename, "Cloudflare").get_item(query)
-    print(f"\n\n\n\n{response}")
+    if response is None:
+        return None
     result = response['result'][0]['results']
     if len(result) == 0:
         return None

db/cloudflare.py

@@ -1,9 +1,7 @@
 import requests, os
-# from dotenv import load_dotenv
 
 class Cloudflare():
     def __init__(self):
-        # load_dotenv()
         self.CF_ACCOOUNT_ID=os.getenv("CF_ACCOOUNT_ID")
         self.DATABASE_ID=os.getenv("DATABASE_ID")
         self.X_Auth_Key=os.getenv("X_Auth_Key")
@@ -27,6 +25,5 @@ class Cloudflare():
         }
         response = requests.post(url, headers=headers, json=input)
         resp_json = response.json()
-
         return resp_json
     

db/tbs_db.py

@@ -20,4 +20,5 @@ class TbsDb():
     
     def execute_query(self, query):
         return self.db_module.execute_query(query)
+
         

db_model/chat.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+
+# 定义消息模型
+class MessageModel(BaseModel):
+    role: str
+    content: str
+
+# 定义聊天模型
+class ChatModel(BaseModel):
+    model: str = None
+    messages: list[MessageModel]

routers/openai_v1.py

@@ -1,64 +1,53 @@
-from fastapi import APIRouter, Request, HTTPException
-# from db import Db
-from dotenv import load_dotenv
-import os, json
+from fastapi import APIRouter, Depends
+import os
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
 
-# 加载.env文件
-load_dotenv()
+from global_state import get
+from db.tbs_db import TbsDb
+from auth import get_current_user
+from db_model.user import UserModel
+from db_model.chat import ChatModel
 
 router = APIRouter()
 
-@router.post("/chat/completions")
-async def chat_completions(request:Request):
-    print('chat_completions')
-    auth_header = request.headers['authorization']
-    user_api_key = auth_header.split()[1]  # 分割字符串并取第二个元素
-    user_id = await get_user_id(user_api_key)
-    if user_id==0:
-        return {'error': {'code': '401', 'message': 'api_key 已过期或验证不正确！'}}
+db_module_filename = f"{get('project_root')}/db/cloudflare.py"
 
-    data = await request.json()
+@router.post("/chat/completions")
+async def chat_completions(chat_model:ChatModel, current_user: UserModel = Depends(get_current_user)):
     try:
-        model = data['model']
+        model = chat_model.model
     except:
         model = ''
-    
-    if model=='':
+
+    if (model=='')or(model is None):
         model = await get_default_model()
-    api_key_info = await get_api_key(model)
 
-    api_key = api_key_info['api_key']
-    group_name = api_key_info['group_name']
-    try:
-        base_url = api_key_info['base_url']
-    except:
-        base_url = ''
+    api_key_info = await get_api_key(model)
+    api_key = api_key_info.get('api_key', '')
+    group_name = api_key_info.get('group_name', '')
+    base_url = api_key_info.get('base_url', '')
 
-    # 下面是google api
-    if group_name=='gemini': 
-        from langchain_google_genai import ChatGoogleGenerativeAI
-        # 初始化模型
+    if group_name=='gemini': # google api，生成 gemini 的 llm
+        from langchain_google_genai import ChatGoogleGenerativeAI    
         llm = ChatGoogleGenerativeAI(
             api_key = api_key,
             model = model,
         )
+    else: # 下面就是 chatgpt 兼容 api
+        from langchain_openai import ChatOpenAI
+        # 初始化 ChatOpenAI 模型
+        llm = ChatOpenAI(
+            model = model,
+            api_key = api_key,
+            base_url = base_url,
+        )
     
-    # 下面就是 chatgpt 兼容 api
-    from langchain_openai import ChatOpenAI
-    # 初始化 ChatOpenAI 模型
-    llm = ChatOpenAI(
-        model = model,
-        api_key = api_key,
-        base_url = base_url,
-    )
-
-    messages = data['messages']
-    rslt = [(item["role"], item["content"]) for item in messages] # 转换为所需的元组列表格式
+    # 生成prompt模板
+    lc_messages = [(message.role, message.content) for message in chat_model.messages]
     parser = StrOutputParser()
-    prompt_template = ChatPromptTemplate.from_messages(rslt)
-
+    prompt_template = ChatPromptTemplate.from_messages(lc_messages)
+ 
     chain = prompt_template | llm | parser
     try:
         result = chain.invoke({})
@@ -66,16 +55,18 @@ async def chat_completions(request:Request):
         return {'error': str(e)}
     return result
 
+# 从数据库获取默认模型
 async def get_default_model():
-    results = Db(os.getenv("DB_PATH")).list_query("SELECT * FROM api_names order by default_order limit 1")
+    query = f"SELECT * FROM api_names order by default_order limit 1"
+    response = TbsDb(db_module_filename, "Cloudflare").get_item(query)
     try:
-        result = results[0]['api_name']
+        result = response['result'][0]['results'][0]['api_name']
     except:
         result = ''
     return result
 
 async def get_api_key(model):
-    sql = f"""
+    query = f"""
 SELECT an.api_name, ak.api_key, an.base_url, ag.group_name
 FROM api_keys ak
 JOIN api_groups ag ON ak.api_group_id = ag.id
@@ -84,16 +75,18 @@ WHERE ak.category='LLM' and an.api_name='{model}' and disabled=0
 ORDER BY ak.last_call_at
 limit 1
 """
-    results = Db(os.getenv("DB_PATH")).list_query(sql)
+    response = TbsDb(db_module_filename, "Cloudflare").get_item(query)
+    # return response
+    # results = Db(os.getenv("DB_PATH")).list_query(sql)
     # return results
     try:
-        result = results[0]
+        result = response['result'][0]['results'][0]
         api_key = result['api_key']
     except:
         api_key = ''
     
-    sql = f"update api_keys set last_call_at=datetime('now') where api_key='{api_key}'"
-    Db(os.getenv("DB_PATH")).update_query(sql)
+    query = f"update api_keys set last_call_at=datetime('now') where api_key='{api_key}'"
+    TbsDb(db_module_filename, "Cloudflare").execute_query(query)
     return result
 
 async def get_user_id(api_key):