fix: throw error if quota exceeded (#717)

.github/workflows/pr-tests.yml

@@ -1,7 +1,7 @@
 name: PR Tests
 
 on:
-  pull_request_target:
+  pull_request:
     branches: [ "main" ]
     paths:
       - .github/workflows/**
@@ -12,7 +12,7 @@ on:
 permissions:
   id-token: write   # This is required for requesting the JWT
   contents: read    # This is required for actions/checkout
-  pull-requests: write 
+  pull-requests: write
   actions: write
 
 env:
@@ -54,7 +54,7 @@ jobs:
           pip install -r requirements.txt
           pip install ruff
           pip install pytest pytest-cov
-        
+
       - name: Lint with Ruff
         run: |
           ruff check --output-format=github .

server/agent/base.py

@@ -2,6 +2,7 @@
 import logging
 from typing import AsyncGenerator, AsyncIterator, Dict, Callable, Optional
 from langchain.agents import AgentExecutor
+from openai import APIError
 from agent.llm import BaseLLMClient
 from petercat_utils.data_class import ChatData, Message
 from langchain.agents.format_scratchpad.openai_tools import (
@@ -152,15 +153,16 @@ async def run_stream_chat(self, input_data: ChatData) -> AsyncIterator[Dict]:
                             "content": content,
                         }
                 elif kind == "on_chat_model_end":
-                    content = event["data"]["output"]["generations"][0][0][
-                        "message"
-                    ].usage_metadata
-                    if content:
-                        yield {
-                            "id": event["run_id"],
-                            "type": "usage",
-                            **content,
-                        }
+                    output = event.get("data", {}).get("output", {})
+                    generations = output.get("generations", [])
+                    if generations and len(generations) > 0:
+                        content = generations[0][0].get("message", {}).get("usage_metadata")
+                        if content:
+                            yield {
+                                "id": event["run_id"],
+                                "type": "usage",
+                                **content,
+                            }
                 elif kind == "on_tool_start":
                     children_value = event["data"].get("input", {})
                     yield {
@@ -213,8 +215,10 @@ async def run_stream_chat(self, input_data: ChatData) -> AsyncIterator[Dict]:
                     }
 
         except Exception as e:
-            res = {"status": "error", "message": str(e)}
-            yield res
+            if isinstance(e, APIError):
+                yield { "status": "error", "message": e.body }
+            else:
+                yield { "status": "error", "message": str(e) }
 
     async def run_chat(self, input_data: ChatData) -> str:
         try:

server/chat/router.py

@@ -36,7 +36,6 @@ def run_qa_chat(
     user: Annotated[User | None, Depends(get_user)] = None,
     bot: Annotated[Bot | None, Depends(get_bot)] = None,
 ):
-    print(f"run_qa_chat: input_data={input_data}, bot={bot}, llm={bot.llm}")
 
     auth_token = (
         Auth.Token(user.access_token) if getattr(user, "access_token", None) else None

server/core/service/user_token_usage.py

@@ -56,6 +56,8 @@ async def record_token_usage(generator: AsyncGenerator[Dict, None]):
               user_token_usage_dao.create(token_usage)
             except Exception as e:
               print(f"An error occurred: {e}")
+          case "error":
+            yield value
           case _:
             yield value
       except Exception as e:

server/insight/service/activity.py

@@ -1,4 +1,3 @@
-import datetime
 import requests
 from typing import List, Dict
 

server/tests/insight/test_activity.py

@@ -35,11 +35,10 @@ def test_get_activity_data_empty(self, mock_get):
 
     @patch("insight.service.activity.requests.get")
     def test_get_activity_data_invalid_json(self, mock_get):
-
         mock_response = MagicMock()
         mock_response.json.side_effect = ValueError("Invalid JSON")
         mock_get.return_value = mock_response
 
         repo_name = "petercat-ai/petercat"
-        with self.assertRaises(ValueError):
-            get_activity_data(repo_name)
+        result = get_activity_data(repo_name)
+        self.assertEqual(result, [])

server/tests/utils/test_insight.py

@@ -1,7 +1,5 @@
 import unittest
 from unittest.mock import patch, Mock
-from collections import defaultdict
-
 from utils.insight import get_data