AgentCoreService/pqai_agent/chat_service.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#

import os
import threading
from typing import List, Dict, Optional
from enum import Enum, auto

import httpx

from pqai_agent import configs
from pqai_agent.logging import logger
import cozepy
from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageType, JWTOAuthApp, JWTAuth
import time
from openai import OpenAI, AsyncOpenAI, http_client

COZE_API_TOKEN = os.getenv("COZE_API_TOKEN")
COZE_CN_BASE_URL = 'https://api.coze.cn'
VOLCENGINE_API_TOKEN = '5e275c38-44fd-415f-abcf-4b59f6377f72'
VOLCENGINE_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"
VOLCENGINE_MODEL_DEEPSEEK_V3 = "deepseek-v3-250324"
VOLCENGINE_MODEL_DOUBAO_PRO_1_5_32K = 'doubao-1-5-pro-32k-250115'
VOLCENGINE_MODEL_DOUBAO_PRO_32K = 'doubao-pro-32k-241215'
VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO = 'doubao-1-5-vision-pro-32k-250115'
DEEPSEEK_API_TOKEN = 'sk-67daad8f424f4854bda7f1fed7ef220b'
DEEPSEEK_BASE_URL = 'https://api.deepseek.com/'
DEEPSEEK_CHAT_MODEL = 'deepseek-chat'
VOLCENGINE_BOT_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3/bots"
VOLCENGINE_BOT_DEEPSEEK_V3_SEARCH = "bot-20250427173459-9h2xp"
OPENAI_API_TOKEN = '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'
OPENAI_BASE_URL = 'https://api.openai.com/v1'
OPENAI_MODEL_GPT_4o = 'gpt-4o'
OPENAI_MODEL_GPT_4o_mini = 'gpt-4o-mini'
OPENROUTER_API_TOKEN = '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'
OPENROUTER_BASE_URL = 'https://openrouter.ai/api/v1/'
OPENROUTER_MODEL_CLAUDE_3_7_SONNET = 'anthropic/claude-3.7-sonnet'
OPENROUTER_MODEL_GEMINI_2_5_PRO = 'google/gemini-2.5-pro'
ALIYUN_API_TOKEN = 'sk-47381479425f4485af7673d3d2fd92b6'
ALIYUN_BASE_URL = 'https://dashscope.aliyuncs.com/compatible-mode/v1'


class ChatServiceType(Enum):
    OPENAI_COMPATIBLE = auto()
    COZE_CHAT = auto()

class ModelPrice:
    EXCHANGE_RATE_TO_CNY = {
        "USD": 7.2,  # Example conversion rate, adjust as needed
    }

    def __init__(self, input_price: float, output_price: float, currency: str = 'CNY'):
        """
        :param input_price: input price for per million tokens
        :param output_price: output price for per million tokens
        """
        self.input_price = input_price
        self.output_price = output_price
        self.currency = currency

    def get_total_cost(self, input_tokens: int, output_tokens: int, convert_to_cny: bool = True) -> float:
        """
        Calculate the total cost based on input and output tokens.
        :param input_tokens: Number of input tokens
        :param output_tokens: Number of output tokens
        :param convert_to_cny: Whether to convert the cost to CNY (default is True)
        :return: Total cost in the specified currency
        """
        total_cost = (self.input_price * input_tokens / 1_000_000) + (self.output_price * output_tokens / 1_000_000)
        if convert_to_cny and self.currency != 'CNY':
            conversion_rate = self.EXCHANGE_RATE_TO_CNY.get(self.currency, 1.0)
            total_cost *= conversion_rate
        return total_cost

    def __repr__(self):
        return f"ModelPrice(input_price={self.input_price}, output_price={self.output_price}, currency={self.currency})"

class OpenAICompatible:
    volcengine_models = [
        VOLCENGINE_MODEL_DOUBAO_PRO_32K,
        VOLCENGINE_MODEL_DOUBAO_PRO_1_5_32K,
        VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO,
        VOLCENGINE_MODEL_DEEPSEEK_V3
    ]
    deepseek_models = [
        DEEPSEEK_CHAT_MODEL,
    ]
    openai_models = [
        OPENAI_MODEL_GPT_4o_mini,
        OPENAI_MODEL_GPT_4o
    ]
    openrouter_models = [
        OPENROUTER_MODEL_CLAUDE_3_7_SONNET,
        OPENROUTER_MODEL_GEMINI_2_5_PRO
    ]

    model_prices = {
        VOLCENGINE_MODEL_DEEPSEEK_V3: ModelPrice(input_price=2, output_price=8),
        VOLCENGINE_MODEL_DOUBAO_PRO_32K: ModelPrice(input_price=0.8, output_price=2),
        VOLCENGINE_MODEL_DOUBAO_PRO_1_5_32K: ModelPrice(input_price=0.8, output_price=2),
        VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO: ModelPrice(input_price=3, output_price=9),
        DEEPSEEK_CHAT_MODEL: ModelPrice(input_price=2, output_price=8),
        OPENAI_MODEL_GPT_4o: ModelPrice(input_price=2.5, output_price=10, currency='USD'),
        OPENAI_MODEL_GPT_4o_mini: ModelPrice(input_price=0.15, output_price=0.6, currency='USD'),
        OPENROUTER_MODEL_CLAUDE_3_7_SONNET: ModelPrice(input_price=3, output_price=15, currency='USD'),
        OPENROUTER_MODEL_GEMINI_2_5_PRO: ModelPrice(input_price=1.25, output_price=10, currency='USD'),
    }

    @staticmethod
    def create_client(model_name, **kwargs) -> OpenAI:
        if model_name in OpenAICompatible.volcengine_models:
            llm_client = OpenAI(api_key=VOLCENGINE_API_TOKEN, base_url=VOLCENGINE_BASE_URL, **kwargs)
        elif model_name in OpenAICompatible.deepseek_models:
            llm_client = OpenAI(api_key=DEEPSEEK_API_TOKEN, base_url=DEEPSEEK_BASE_URL, **kwargs)
        elif model_name in OpenAICompatible.openai_models:
            socks_conf = configs.get().get('system', {}).get('outside_proxy', {}).get('socks5', {})
            if socks_conf:
                http_client = httpx.Client(
                    timeout=httpx.Timeout(600, connect=5.0),
                    proxy=f"socks5://{socks_conf['hostname']}:{socks_conf['port']}"
                )
                kwargs['http_client'] = http_client
            llm_client = OpenAI(api_key=OPENAI_API_TOKEN, base_url=OPENAI_BASE_URL, **kwargs)
        elif model_name in OpenAICompatible.openrouter_models:
            llm_client = OpenAI(api_key=OPENROUTER_API_TOKEN, base_url=OPENROUTER_BASE_URL, **kwargs)
        else:
            raise Exception("Unsupported model: %s" % model_name)
        return llm_client

    @staticmethod
    def get_price(model_name: str) -> ModelPrice:
        """
        Get the price for a given model.
        :param model_name: Name of the model
        :return: ModelPrice object containing input and output prices
        """
        if model_name not in OpenAICompatible.model_prices:
            raise ValueError(f"Model {model_name} not found in price list.")
        return OpenAICompatible.model_prices[model_name]

    @staticmethod
    def calculate_cost(model_name: str, input_tokens: int, output_tokens: int, convert_to_cny: bool = True) -> float:
        """
        Calculate the cost for a given model based on input and output tokens.
        :param model_name: Name of the model
        :param input_tokens: Number of input tokens
        :param output_tokens: Number of output tokens
        :param convert_to_cny: Whether to convert the cost to CNY (default is True)
        :return: Total cost in the model's currency
        """
        if model_name not in OpenAICompatible.model_prices:
            raise ValueError(f"Model {model_name} not found in price list.")
        price = OpenAICompatible.model_prices[model_name]
        return price.get_total_cost(input_tokens, output_tokens, convert_to_cny)

class CrossAccountJWTOAuthApp(JWTOAuthApp):
    def __init__(self, account_id: str, client_id: str, private_key: str, public_key_id: str, base_url):
        self.account_id = account_id
        super().__init__(client_id, private_key, public_key_id, base_url)

    def get_access_token(
            self, ttl: int = 900, scope: Optional[cozepy.Scope] = None, session_name: Optional[str] = None
    ) -> cozepy.OAuthToken:
        jwt_token = self._gen_jwt(self._public_key_id, self._private_key, 3600, session_name)
        url = f"{self._base_url}/api/permission/oauth2/account/{self.account_id}/token"
        headers = {"Authorization": f"Bearer {jwt_token}"}
        body = {
            "duration_seconds": ttl,
            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
            "scope": scope.model_dump() if scope else None,
        }
        return self._requester.request("post", url, False, cozepy.OAuthToken, headers=headers, body=body)

class CozeChat:
    def __init__(self, base_url: str, auth_token: Optional[str] = None, auth_app: Optional[JWTOAuthApp] = None):
        if not auth_token and not auth_app:
            raise ValueError("Either auth_token or auth_app must be provided.")
        self.thread = None
        self.thread_running = False
        self.last_token_fresh = 0
        if auth_token:
            self.coze = Coze(auth=TokenAuth(auth_token), base_url=base_url)
        else:
            self.auth_app = auth_app
            oauth_token = auth_app.get_access_token(ttl=12*3600)
            self.last_token_fresh = time.time()
            self.coze = Coze(auth=JWTAuth(oauth_app=auth_app), base_url=base_url)
            self.setup_token_refresh()

    def create(self, bot_id: str, user_id: str, messages: List, custom_variables: Dict):
        response = self.coze.chat.create_and_poll(
            bot_id=bot_id, user_id=user_id, additional_messages=messages,
            custom_variables=custom_variables)
        logger.debug("Coze response size: {}".format(len(response.messages)))
        if response.chat.status != ChatStatus.COMPLETED:
            logger.error("Coze chat not completed: {}".format(response.chat.status))
            return None
        final_response = None
        for message in response.messages:
            if message.type == MessageType.ANSWER:
                final_response = message.content
        return final_response

    def setup_token_refresh(self):
        self.thread = threading.Thread(target=self.refresh_token_loop)
        self.thread.start()
        self.thread_running = True

    def refresh_token_loop(self):
        while self.thread_running:
            if time.time() - self.last_token_fresh < 11*3600:
                time.sleep(1)
                continue
            if self.auth_app:
                self.auth_app.get_access_token(ttl=12*3600)
                self.last_token_fresh = time.time()

    def __del__(self):
        self.thread_running = False


    @staticmethod
    def get_oauth_app(client_id, private_key_path, public_key_id, base_url=None, account_id=None) -> JWTOAuthApp:
        if not base_url:
            base_url = COZE_CN_BASE_URL
        with open(private_key_path, "r") as f:
            private_key = f.read()
        if not account_id:
            jwt_oauth_app = JWTOAuthApp(
                client_id=str(client_id),
                private_key=private_key,
                public_key_id=public_key_id,
                base_url=base_url,
            )
        else:
            jwt_oauth_app = CrossAccountJWTOAuthApp(
                account_id=account_id,
                client_id=str(client_id),
                private_key=private_key,
                public_key_id=public_key_id,
                base_url=base_url,
            )
        return jwt_oauth_app

if __name__ == '__main__':
    # Init the Coze client through the access_token.
    coze = Coze(auth=TokenAuth(token=COZE_API_TOKEN), base_url=COZE_CN_BASE_URL)

    # Create a bot instance in Coze, copy the last number from the web link as the bot's ID.
    bot_id = "7491250992952999973"
    # The user id identifies the identity of a user. Developers can use a custom business ID
    # or a random string.
    user_id = "dev_user"

    chat = coze.chat.create_and_poll(
        bot_id=bot_id,
        user_id=user_id,
        additional_messages=[Message.build_user_question_text("钱塘江边 樱花开得不错，推荐一个视频吧")],
        custom_variables={
            'agent_name': '芳华',
            'agent_age': '25',
            'agent_region': '北京',
            'name': '李明',
            'preferred_nickname': '李叔',
            'age': '70',
            'last_interaction_interval': '12',
            'current_time_period': '上午',
            'if_first_interaction': 'False',
            'if_active_greeting': 'False'
        }
    )

    for message in chat.messages:
        print(message, flush=True)

    if chat.chat.status == ChatStatus.COMPLETED:
        print("token usage:", chat.chat.usage.token_count)