diff --git a/AIGC/AICore.py b/AIGC/AICore.py new file mode 100644 index 0000000..378e2ca --- /dev/null +++ b/AIGC/AICore.py @@ -0,0 +1,21 @@ +import requests +from ollama import Client, ResponseError +import tiktoken + +class AICore: + modelMaxTokens = 128000 + # 初始化 DeepSeek 使用的 Tokenizer (cl100k_base) + encoder = tiktoken.get_encoding("cl100k_base") + + + + def __init__(self, model): + #初始化ollama客户端 + ollamaClient = Client(host='http://localhost:11434', headers={'x-some-header': 'some-value'}) + response = ollamaClient.show(model) + modelMaxTokens = response.modelinfo['qwen2.context_length'] + + def getPromptToken(self, prompt)-> int: + tokens = self.encoder.encode(prompt) + return len(tokens) + \ No newline at end of file diff --git a/AIGC/main.py b/AIGC/main.py index 9c50f78..1fdb06c 100644 --- a/AIGC/main.py +++ b/AIGC/main.py @@ -11,8 +11,9 @@ from fastapi import FastAPI, Request, HTTPException, WebSocket, WebSocketDisconn from fastapi.websockets import WebSocketState from h11 import ConnectionClosed import uvicorn +from AICore import AICore from Utils.AIGCLog import AIGCLog -from ollama import Client, ResponseError + app = FastAPI(title = "AI 通信服务") logger = AIGCLog(name = "AIGC", log_file = "aigc.log") @@ -27,8 +28,8 @@ logger.log(logging.INFO, f"使用的模型是 {args.model}") maxAIRegerateCount = 5 lastPrompt = "" -#初始化ollama客户端 -ollamaClient = Client(host='http://localhost:11434') +aicore = AICore(args.model) + async def heartbeat(websocket: WebSocket): pass @@ -150,6 +151,7 @@ async def generateAIChat(promptStr: str, websocket: WebSocket| None = None): {"role": "system", "content": promptStr} ] try: + # response = ollamaClient.chat( # model = args.model, # stream = False, @@ -235,6 +237,7 @@ if __name__ == "__main__": server_thread.start() # Test + aicore.getPromptToken("测试功能") asyncio.run( generateAIChat(promptStr = f""" #你是一个游戏NPC对话生成器。请严格按以下要求生成两个角色的日常对话 diff --git a/AIGC/requirements.txt b/AIGC/requirements.txt index 909a897..e961b7d 100644 --- a/AIGC/requirements.txt +++ b/AIGC/requirements.txt @@ -1,3 +1,4 @@ uvicorn[standard] fastapi -ollama \ No newline at end of file +ollama +tiktoken \ No newline at end of file