Project02/AIGC/AICore.py

import logging
from typing import Tuple
import requests
from ollama import Client, ResponseError
import tiktoken
import random
from Utils.AIGCLog import AIGCLog

class AICore:
    modelMaxTokens = 128000
    # 初始化 DeepSeek 使用的 Tokenizer (cl100k_base)
    encoder = tiktoken.get_encoding("cl100k_base")
    logger = AIGCLog(name = "AIGC", log_file = "aigc.log")


    def __init__(self, model):
         #初始化ollama客户端
        self.ollamaClient = Client(host='http://localhost:11434', headers={'x-some-header': 'some-value'})
        self.modelName = model
        response = self.ollamaClient.show(model)
        modelMaxTokens = response.modelinfo['qwen2.context_length']
        
    def getPromptToken(self, prompt)-> int: 
        tokens = self.encoder.encode(prompt)
        return len(tokens)
    
    def generateAI(self, promptStr: str) -> Tuple[bool, str]:
        try:
            response = self.ollamaClient.generate(
                model = self.modelName,
                stream = False,
                prompt = promptStr,
                options={
                    "temperature": random.uniform(1.0, 1.5),
                    "repeat_penalty": 1.2,  # 抑制重复
                    "top_p": random.uniform(0.7, 0.95),
                    "num_ctx": 4096,           # 上下文长度
                }
            )
            return True, response.response
        except ResponseError as e:
            if e.status_code == 503:
                print("🔄 服务不可用，5秒后重试...")    
            return False,"ollama 服务不可用"
        except Exception as e:
            print(f"🔥 未预料错误: {str(e)}")
            return False, "未预料错误"