feat(copilot): Premium request should be counted when user initial chat

Author: daltonnyx

AgentCrew/modules/a2a/server.py

@@ -23,6 +23,10 @@
     InvalidRequestError,
     JSONParseError,
     InternalError,
+    SendMessageRequest,
+    SendStreamingMessageRequest,
+    GetTaskRequest,
+    CancelTaskRequest,
 )
 
 logger = logging.getLogger(__name__)
@@ -160,13 +164,17 @@ async def process_jsonrpc_request(request: Request):
                 method = json_rpc_request.root.method
                 logger.debug(f"Processing method: {method}")
 
-                if method == "message/send":
+                if method == "message/send" and isinstance(
+                    json_rpc_request.root, SendMessageRequest
+                ):
                     logger.debug("Handling message/send request")
                     result = await task_manager.on_send_message(json_rpc_request.root)
                     logger.debug(f"message/send result: {result}")
                     return JSONResponse(result.model_dump(exclude_none=True))
 
-                elif method == "message/stream":
+                elif method == "message/stream" and isinstance(
+                    json_rpc_request.root, SendStreamingMessageRequest
+                ):
                     result_stream = task_manager.on_send_message_streaming(
                         json_rpc_request.root
                     )
@@ -182,13 +190,17 @@ async def event_generator():
 
                     return EventSourceResponse(event_generator())
 
-                elif method == "tasks/send":
+                elif method == "tasks/send" and isinstance(
+                    json_rpc_request.root, SendMessageRequest
+                ):
                     logger.debug("Handling legacy tasks/send request")
                     result = await task_manager.on_send_task(json_rpc_request.root)
                     logger.debug(f"tasks/send result: {result}")
                     return JSONResponse(result.model_dump(exclude_none=True))
 
-                elif method == "tasks/sendSubscribe":
+                elif method == "tasks/sendSubscribe" and isinstance(
+                    json_rpc_request.root, SendStreamingMessageRequest
+                ):
                     result_stream = task_manager.on_send_task_subscribe(
                         json_rpc_request.root
                     )
@@ -204,11 +216,15 @@ async def event_generator():
 
                     return EventSourceResponse(event_generator())
 
-                elif method == "tasks/get":
+                elif method == "tasks/get" and isinstance(
+                    json_rpc_request.root, GetTaskRequest
+                ):
                     result = await task_manager.on_get_task(json_rpc_request.root)
                     return JSONResponse(result.model_dump(exclude_none=True))
 
-                elif method == "tasks/cancel":
+                elif method == "tasks/cancel" and isinstance(
+                    json_rpc_request.root, CancelTaskRequest
+                ):
                     result = await task_manager.on_cancel_task(json_rpc_request.root)
                     return JSONResponse(result.model_dump(exclude_none=True))
 

AgentCrew/modules/custom_llm/github_copilot_service.py

@@ -1,7 +1,9 @@
+from uuid import uuid4
 from AgentCrew.modules.custom_llm import CustomLLMService
 import os
 from dotenv import load_dotenv
 from AgentCrew.modules import logger
+from datetime import datetime
 
 
 class GithubCopilotService(CustomLLMService):
@@ -27,4 +29,70 @@ def __init__(self):
         self.current_output_tokens = 0
         self.temperature = 0.6
         self._is_thinking = False
+        # self._interaction_id = None
         logger.info("Initialized Github Copilot Service")
+
+    def github_copilot_token_to_open_ai_key(self, copilot_api_key):
+        """
+        Convert GitHub Copilot token to OpenAI key format.
+
+        Args:
+            copilot_api_key: The GitHub Copilot token
+
+        Returns:
+            Updated OpenAI compatible token
+        """
+        openai_api_key = self.client.api_key
+
+        if openai_api_key.startswith("ghu") or int(
+            dict(x.split("=") for x in openai_api_key.split(";"))["exp"]
+        ) < int(datetime.now().timestamp()):
+            import requests
+
+            headers = {
+                "Authorization": f"Bearer {copilot_api_key}",
+                "Content-Type": "application/json",
+            }
+            if self.extra_headers:
+                headers.update(self.extra_headers)
+            res = requests.get(
+                "https://api.github.com/copilot_internal/v2/token", headers=headers
+            )
+            self.client.api_key = res.json()["token"]
+
+    async def process_message(self, prompt: str, temperature: float = 0) -> str:
+        # Check if using GitHub Copilot
+        if self.base_url:
+            from urllib.parse import urlparse
+
+            parsed_url = urlparse(self.base_url)
+            host = parsed_url.hostname
+            if host and host.endswith(".githubcopilot.com"):
+                self.base_url = self.base_url.rstrip("/")
+                self.github_copilot_token_to_open_ai_key(self.api_key)
+        return await super().process_message(prompt, temperature)
+
+    async def stream_assistant_response(self, messages):
+        """Stream the assistant's response with tool support."""
+        # Check if using GitHub Copilot
+        if self.base_url:
+            from urllib.parse import urlparse
+
+            parsed_url = urlparse(self.base_url)
+            host = parsed_url.hostname
+            if host and host.endswith(".githubcopilot.com"):
+                self.base_url = self.base_url.rstrip("/")
+                self.github_copilot_token_to_open_ai_key(self.api_key)
+                # if len([m for m in messages if m.get("role") == "assistant"]) == 0:
+                #     self._interaction_id = str(uuid4())
+                if self.extra_headers:
+                    self.extra_headers["X-Initiator"] = (
+                        "user"
+                        if messages[-1].get("role", "assistant") == "user"
+                        else "agent"
+                    )
+                    self.extra_headers["X-Request-Id"] = str(uuid4())
+                    # if self._interaction_id:
+                    #     self.extra_headers["X-Interaction-Id"] = self._interaction_id
+
+        return await super().stream_assistant_response(messages)

AgentCrew/modules/custom_llm/service.py

@@ -4,7 +4,6 @@
 from mcp.types import ImageContent, TextContent
 from typing import Dict, Any, List, Optional, Tuple
 import json
-from datetime import datetime
 from AgentCrew.modules import logger
 
 
@@ -75,16 +74,6 @@ def format_tool_result(
 
     async def process_message(self, prompt: str, temperature: float = 0) -> str:
         try:
-            # Check if using GitHub Copilot
-            if self.base_url:
-                from urllib.parse import urlparse
-
-                parsed_url = urlparse(self.base_url)
-                host = parsed_url.hostname
-                if host and host.endswith(".githubcopilot.com"):
-                    self.base_url = self.base_url.rstrip("/")
-                    self.github_copilot_token_to_open_ai_key(self.api_key)
-
             response = await self.client.chat.completions.create(
                 model=self.model,
                 max_tokens=3000,
@@ -132,15 +121,6 @@ async def process_message(self, prompt: str, temperature: float = 0) -> str:
 
     async def stream_assistant_response(self, messages):
         """Stream the assistant's response with tool support."""
-        # Check if using GitHub Copilot
-        if self.base_url:
-            from urllib.parse import urlparse
-
-            parsed_url = urlparse(self.base_url)
-            host = parsed_url.hostname
-            if host and host.endswith(".githubcopilot.com"):
-                self.base_url = self.base_url.rstrip("/")
-                self.github_copilot_token_to_open_ai_key(self.api_key)
 
         stream_params = {
             "model": self.model,
@@ -203,34 +183,6 @@ def process_stream_chunk(
         else:
             return self._process_non_stream_chunk(chunk, assistant_response, tool_uses)
 
-    def github_copilot_token_to_open_ai_key(self, copilot_api_key):
-        """
-        Convert GitHub Copilot token to OpenAI key format.
-
-        Args:
-            copilot_api_key: The GitHub Copilot token
-
-        Returns:
-            Updated OpenAI compatible token
-        """
-        openai_api_key = self.client.api_key
-
-        if openai_api_key.startswith("ghu") or int(
-            dict(x.split("=") for x in openai_api_key.split(";"))["exp"]
-        ) < int(datetime.now().timestamp()):
-            import requests
-
-            headers = {
-                "Authorization": f"Bearer {copilot_api_key}",
-                "Content-Type": "application/json",
-            }
-            if self.extra_headers:
-                headers.update(self.extra_headers)
-            res = requests.get(
-                "https://api.github.com/copilot_internal/v2/token", headers=headers
-            )
-            self.client.api_key = res.json()["token"]
-
     def _process_non_stream_chunk(
         self, chunk, assistant_response, tool_uses
     ) -> Tuple[str, List[Dict], int, int, Optional[str], Optional[tuple]]: