Update src/analyzer/openai_analyzer.py

src/analyzer/openai_analyzer.py

@@ -1,46 +1,31 @@
+# token_utils.py
+import tiktoken
 
-from typing import List
-from openai import OpenAI
-
-from src.core.types import DEFAULT_OPENAI_ANALYZER, DEFAULT_SYSTEM_PROMPT, DEFAULT_USER_PROMPT
-
-from src.models.analyzer_models import AnalyzerResult
-from src.models.scrape_models import ScrapeResult
-
-from src.core.interface.analyzer_interface import AnalyzerInterface
-
-
-class OpenaiAnalyzer(AnalyzerInterface):
-    def __init__(self, api_key, model_name = DEFAULT_OPENAI_ANALYZER):
-       self.client = OpenAI(api_key=api_key)
-       self.model_name = model_name
-
-    def analyze_search_result(self, query: str, search_results: List[ScrapeResult]) -> AnalyzerResult:
-        """
-        Analyzes the provided search results based on the given query.
-        Args:
-            query (str): The search query string.
-            search_results (List[ScrapeResult]): A list of search results to be analyzed.
-        Returns:
-            AnalyzerResult: The result of the analysis.
-        Raises:
-            NotImplementedError: If the method is not implemented by a subclass.
-        """
+class TokenCounter:
+    """Utility class for counting and managing tokens."""
+    
+    @staticmethod
+    def get_encoder(model_name: str):
+        """Get appropriate encoder for model."""
         try:
-            user_prompt = DEFAULT_USER_PROMPT.replace("query", query).replace("scrape_results", f"{search_results}")
-            completion = self.client.beta.chat.completions.parse(model=self.model_name,
-                                                                 messages=[
-                                                                     {
-                                                                         "role": "system",
-                                                                         "content": DEFAULT_SYSTEM_PROMPT
-                                                                     },
-                                                                     {
-                                                                         "role": "user",
-                                                                         "content": user_prompt
-                                                                     }
-                                                                 ],
-                                                                 response_format=AnalyzerResult)
-            response = completion.choices[0].message.parsed
-            return response
-        except Exception as e:
-            raise Exception(f"Error while analyzing search result: {str(e)}")
+            return tiktoken.encoding_for_model(model_name)
+        except:
+            return tiktoken.get_encoding("cl100k_base")  # Default for newer models
+    
+    @staticmethod
+    def estimate_prompt_tokens(model_name: str, system_prompt: str, user_prompt: str) -> int:
+        """Estimate total tokens in a prompt."""
+        encoder = TokenCounter.get_encoder(model_name)
+        return len(encoder.encode(system_prompt)) + len(encoder.encode(user_prompt))
+    
+    @staticmethod
+    def get_model_context_limit(model_name: str) -> int:
+        """Get context limit for a given model."""
+        limits = {
+            "gpt-3.5-turbo": 16385,
+            "gpt-4": 8192,
+            "gpt-4-turbo": 128000,
+            "gpt-4o": 128000,
+            "gpt-4o-mini": 128000,
+        }
+        return limits.get(model_name, 8192)  # Default to 8192 if unknown