Prechádzať zdrojové kódy

Update api_server: display model price

StrayWarrior 2 dní pred
rodič
commit
1d30539ffc

+ 5 - 0
pqai_agent/chat_service.py

@@ -74,6 +74,11 @@ class ModelPrice:
             total_cost *= conversion_rate
         return total_cost
 
+    def get_cny_brief(self) -> str:
+        input_price = self.input_price * self.EXCHANGE_RATE_TO_CNY.get(self.currency, 1.0)
+        output_price = self.output_price * self.EXCHANGE_RATE_TO_CNY.get(self.currency, 1.0)
+        return f"{input_price:.0f}/{output_price:.0f}"
+
     def __repr__(self):
         return f"ModelPrice(input_price={self.input_price}, output_price={self.output_price}, currency={self.currency})"
 

+ 2 - 1
pqai_agent_server/api_server.py

@@ -13,6 +13,7 @@ from sqlalchemy.orm import sessionmaker
 from pqai_agent import configs
 
 from pqai_agent import chat_service, prompt_templates
+from pqai_agent.chat_service import OpenAICompatible
 from pqai_agent.logging import logger, setup_root_logger
 from pqai_agent.toolkit import global_tool_map
 from pqai_agent.data_models.agent_configuration import AgentConfiguration
@@ -106,7 +107,7 @@ def list_models():
         {
             'model_type': 'openai_compatible',
             'model_name': model_name,
-            'display_name': model_display_name
+            'display_name': f"{model_display_name} ({OpenAICompatible.get_price(model_name).get_cny_brief()})"
         }
         for model_display_name, model_name in models.items()
     ]