Further simplify ChatModel:

- Remove provider-specific model files (LiteLLM handles it) - Remove token count tracking (not needed) - Update tests to match simpler response format
wandb · Dec 26, 2024 · 9171818 · 9171818
1 parent ef7f29b
commit 9171818
Show file tree

Hide file tree

Showing 12 changed files with 2 additions and 1,184 deletions.
diff --git a/src/wandbot/chat/chat_model.py b/src/wandbot/chat/chat_model.py
@@ -123,19 +123,13 @@ def generate_response(
 
             return {
                 "content": response.choices[0].message.content,
-                "total_tokens": response.usage.total_tokens,
-                "prompt_tokens": response.usage.prompt_tokens,
-                "completion_tokens": response.usage.completion_tokens,
                 "error": None,
                 "model_used": response.model
             }
 
         except Exception as e:
             return {
                 "content": "",
-                "total_tokens": 0,
-                "prompt_tokens": 0,
-                "completion_tokens": 0,
                 "error": {
                     "type": type(e).__name__,
                     "message": str(e)

diff --git a/src/wandbot/chat/models/__init__.py b/src/wandbot/chat/models/__init__.py
diff --git a/src/wandbot/chat/models/anthropic_model.py b/src/wandbot/chat/models/anthropic_model.py
diff --git a/src/wandbot/chat/models/base.py b/src/wandbot/chat/models/base.py
diff --git a/src/wandbot/chat/models/gemini_model.py b/src/wandbot/chat/models/gemini_model.py