refactor and deps (#2184)

application/llm/docsgpt_provider.py

@@ -1,75 +1,19 @@
-import json
-
-from openai import OpenAI
-
 from application.core.settings import settings
-from application.llm.base import BaseLLM
+from application.llm.openai import OpenAILLM
 
+DOCSGPT_API_KEY = "sk-docsgpt-public"
+DOCSGPT_BASE_URL = "https://oai.arc53.com"
+DOCSGPT_MODEL = "docsgpt"
 
-class DocsGPTAPILLM(BaseLLM):
-
-    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
-
-        super().__init__(*args, **kwargs)
-        self.api_key = "sk-docsgpt-public"
-        self.client = OpenAI(api_key=self.api_key, base_url="https://oai.arc53.com")
-        self.user_api_key = user_api_key
-
-    def _clean_messages_openai(self, messages):
-        cleaned_messages = []
-        for message in messages:
-            role = message.get("role")
-            content = message.get("content")
-
-            if role == "model":
-                role = "assistant"
-            if role and content is not None:
-                if isinstance(content, str):
-                    cleaned_messages.append({"role": role, "content": content})
-                elif isinstance(content, list):
-                    for item in content:
-                        if "text" in item:
-                            cleaned_messages.append(
-                                {"role": role, "content": item["text"]}
-                            )
-                        elif "function_call" in item:
-                            cleaned_args = self._remove_null_values(
-                                item["function_call"]["args"]
-                            )
-                            tool_call = {
-                                "id": item["function_call"]["call_id"],
-                                "type": "function",
-                                "function": {
-                                    "name": item["function_call"]["name"],
-                                    "arguments": json.dumps(cleaned_args),
-                                },
-                            }
-                            cleaned_messages.append(
-                                {
-                                    "role": "assistant",
-                                    "content": None,
-                                    "tool_calls": [tool_call],
-                                }
-                            )
-                        elif "function_response" in item:
-                            cleaned_messages.append(
-                                {
-                                    "role": "tool",
-                                    "tool_call_id": item["function_response"][
-                                        "call_id"
-                                    ],
-                                    "content": json.dumps(
-                                        item["function_response"]["response"]["result"]
-                                    ),
-                                }
-                            )
-                        else:
-                            raise ValueError(
-                                f"Unexpected content dictionary format: {item}"
-                            )
-                else:
-                    raise ValueError(f"Unexpected content type: {type(content)}")
-        return cleaned_messages
+class DocsGPTAPILLM(OpenAILLM):
+    def __init__(self, api_key=None, user_api_key=None, base_url=None, *args, **kwargs):
+        super().__init__(
+            api_key=DOCSGPT_API_KEY,
+            user_api_key=user_api_key,
+            base_url=DOCSGPT_BASE_URL,
+            *args,
+            **kwargs,
+        )
 
     def _raw_gen(
         self,
@@ -79,23 +23,19 @@ class DocsGPTAPILLM(BaseLLM):
         stream=False,
         tools=None,
         engine=settings.AZURE_DEPLOYMENT_NAME,
+        response_format=None,
         **kwargs,
     ):
-        messages = self._clean_messages_openai(messages)
-        if tools:
-            response = self.client.chat.completions.create(
-                model="docsgpt",
-                messages=messages,
-                stream=stream,
-                tools=tools,
-                **kwargs,
-            )
-            return response.choices[0]
-        else:
-            response = self.client.chat.completions.create(
-                model="docsgpt", messages=messages, stream=stream, **kwargs
-            )
-            return response.choices[0].message.content
+        return super()._raw_gen(
+            baseself,
+            DOCSGPT_MODEL,
+            messages,
+            stream=stream,
+            tools=tools,
+            engine=engine,
+            response_format=response_format,
+            **kwargs,
+        )
 
     def _raw_gen_stream(
         self,
@@ -105,34 +45,16 @@ class DocsGPTAPILLM(BaseLLM):
         stream=True,
         tools=None,
         engine=settings.AZURE_DEPLOYMENT_NAME,
+        response_format=None,
         **kwargs,
     ):
-        messages = self._clean_messages_openai(messages)
-        if tools:
-            response = self.client.chat.completions.create(
-                model="docsgpt",
-                messages=messages,
-                stream=stream,
-                tools=tools,
-                **kwargs,
-            )
-        else:
-            response = self.client.chat.completions.create(
-                model="docsgpt", messages=messages, stream=stream, **kwargs
-            )
-        try:
-            for line in response:
-                if (
-                    len(line.choices) > 0
-                    and line.choices[0].delta.content is not None
-                    and len(line.choices[0].delta.content) > 0
-                ):
-                    yield line.choices[0].delta.content
-                elif len(line.choices) > 0:
-                    yield line.choices[0]
-        finally:
-            if hasattr(response, "close"):
-                response.close()
-
-    def _supports_tools(self):
-        return True
+        return super()._raw_gen_stream(
+            baseself,
+            DOCSGPT_MODEL,
+            messages,
+            stream=stream,
+            tools=tools,
+            engine=engine,
+            response_format=response_format,
+            **kwargs,
+        )

application/llm/groq.py

@@ -1,37 +1,15 @@
-from openai import OpenAI
-
 from application.core.settings import settings
-from application.llm.base import BaseLLM
+from application.llm.openai import OpenAILLM
+
+GROQ_BASE_URL = "https://api.groq.com/openai/v1"
 
 
-class GroqLLM(BaseLLM):
-    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
-
-        super().__init__(*args, **kwargs)
-        self.api_key = api_key or settings.GROQ_API_KEY or settings.API_KEY
-        self.user_api_key = user_api_key
-        self.client = OpenAI(
-            api_key=self.api_key, base_url="https://api.groq.com/openai/v1"
+class GroqLLM(OpenAILLM):
+    def __init__(self, api_key=None, user_api_key=None, base_url=None, *args, **kwargs):
+        super().__init__(
+            api_key=api_key or settings.GROQ_API_KEY or settings.API_KEY,
+            user_api_key=user_api_key,
+            base_url=base_url or GROQ_BASE_URL,
+            *args,
+            **kwargs,
         )
-
-    def _raw_gen(self, baseself, model, messages, stream=False, tools=None, **kwargs):
-        if tools:
-            response = self.client.chat.completions.create(
-                model=model, messages=messages, stream=stream, tools=tools, **kwargs
-            )
-            return response.choices[0]
-        else:
-            response = self.client.chat.completions.create(
-                model=model, messages=messages, stream=stream, **kwargs
-            )
-            return response.choices[0].message.content
-
-    def _raw_gen_stream(
-        self, baseself, model, messages, stream=True, tools=None, **kwargs
-    ):
-        response = self.client.chat.completions.create(
-            model=model, messages=messages, stream=stream, **kwargs
-        )
-        for line in response:
-            if line.choices[0].delta.content is not None:
-                yield line.choices[0].delta.content

application/llm/handlers/base.py

@@ -833,7 +833,10 @@ class LLMHandler(ABC):
                         if call.name:
                             existing.name = call.name
                         if call.arguments:
-                            existing.arguments += call.arguments
+                            if existing.arguments is None:
+                                existing.arguments = call.arguments
+                            else:
+                                existing.arguments += call.arguments
                         # Preserve thought_signature for Google Gemini 3 models
                         if call.thought_signature:
                             existing.thought_signature = call.thought_signature

application/llm/huggingface.py

@@ -1,68 +0,0 @@
-from application.llm.base import BaseLLM
-
-
-class HuggingFaceLLM(BaseLLM):
-
-    def __init__(
-        self,
-        api_key=None,
-        user_api_key=None,
-        llm_name="Arc53/DocsGPT-7B",
-        q=False,
-        *args,
-        **kwargs,
-    ):
-        global hf
-
-        from langchain.llms import HuggingFacePipeline
-
-        if q:
-            import torch
-            from transformers import (
-                AutoModelForCausalLM,
-                AutoTokenizer,
-                pipeline,
-                BitsAndBytesConfig,
-            )
-
-            tokenizer = AutoTokenizer.from_pretrained(llm_name)
-            bnb_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_use_double_quant=True,
-                bnb_4bit_quant_type="nf4",
-                bnb_4bit_compute_dtype=torch.bfloat16,
-            )
-            model = AutoModelForCausalLM.from_pretrained(
-                llm_name, quantization_config=bnb_config
-            )
-        else:
-            from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-
-            tokenizer = AutoTokenizer.from_pretrained(llm_name)
-            model = AutoModelForCausalLM.from_pretrained(llm_name)
-
-        super().__init__(*args, **kwargs)
-        self.api_key = api_key
-        self.user_api_key = user_api_key
-        pipe = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            max_new_tokens=2000,
-            device_map="auto",
-            eos_token_id=tokenizer.eos_token_id,
-        )
-        hf = HuggingFacePipeline(pipeline=pipe)
-
-    def _raw_gen(self, baseself, model, messages, stream=False, **kwargs):
-        context = messages[0]["content"]
-        user_question = messages[-1]["content"]
-        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
-
-        result = hf(prompt)
-
-        return result.content
-
-    def _raw_gen_stream(self, baseself, model, messages, stream=True, **kwargs):
-
-        raise NotImplementedError("HuggingFaceLLM Streaming is not implemented yet.")

application/llm/llm_creator.py

@@ -4,7 +4,6 @@ from application.llm.anthropic import AnthropicLLM
 from application.llm.docsgpt_provider import DocsGPTAPILLM
 from application.llm.google_ai import GoogleLLM
 from application.llm.groq import GroqLLM
-from application.llm.huggingface import HuggingFaceLLM
 from application.llm.llama_cpp import LlamaCpp
 from application.llm.novita import NovitaLLM
 from application.llm.openai import AzureOpenAILLM, OpenAILLM
@@ -19,7 +18,6 @@ class LLMCreator:
         "openai": OpenAILLM,
         "azure_openai": AzureOpenAILLM,
         "sagemaker": SagemakerAPILLM,
-        "huggingface": HuggingFaceLLM,
         "llama.cpp": LlamaCpp,
         "anthropic": AnthropicLLM,
         "docsgpt": DocsGPTAPILLM,

application/llm/novita.py

@@ -1,32 +1,15 @@
-from application.llm.base import BaseLLM
-from openai import OpenAI
+from application.core.settings import settings
+from application.llm.openai import OpenAILLM
+
+NOVITA_BASE_URL = "https://api.novita.ai/v3/openai"
 
 
-class NovitaLLM(BaseLLM):
-    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.client = OpenAI(api_key=api_key, base_url="https://api.novita.ai/v3/openai")
-        self.api_key = api_key
-        self.user_api_key = user_api_key
-
-    def _raw_gen(self, baseself, model, messages, stream=False, tools=None, **kwargs):
-        if tools:
-            response = self.client.chat.completions.create(
-                model=model, messages=messages, stream=stream, tools=tools, **kwargs
-            )
-            return response.choices[0]
-        else:
-            response = self.client.chat.completions.create(
-                model=model, messages=messages, stream=stream, **kwargs
-            )
-            return response.choices[0].message.content
-
-    def _raw_gen_stream(
-        self, baseself, model, messages, stream=True, tools=None, **kwargs
-    ):
-        response = self.client.chat.completions.create(
-            model=model, messages=messages, stream=stream, **kwargs
+class NovitaLLM(OpenAILLM):
+    def __init__(self, api_key=None, user_api_key=None, base_url=None, *args, **kwargs):
+        super().__init__(
+            api_key=api_key or settings.API_KEY,
+            user_api_key=user_api_key,
+            base_url=base_url or NOVITA_BASE_URL,
+            *args,
+            **kwargs,
         )
-        for line in response:
-            if line.choices[0].delta.content is not None:
-                yield line.choices[0].delta.content

application/llm/openai.py

@@ -127,6 +127,7 @@ class OpenAILLM(BaseLLM):
         **kwargs,
     ):
         messages = self._clean_messages_openai(messages)
+        logging.info(f"Cleaned messages: {messages}")
 
         # Convert max_tokens to max_completion_tokens for newer models
         if "max_tokens" in kwargs:
@@ -144,7 +145,7 @@ class OpenAILLM(BaseLLM):
         if response_format:
             request_params["response_format"] = response_format
         response = self.client.chat.completions.create(**request_params)
-
+        logging.info(f"OpenAI response: {response}")
         if tools:
             return response.choices[0]
         else:
@@ -162,6 +163,7 @@ class OpenAILLM(BaseLLM):
         **kwargs,
     ):
         messages = self._clean_messages_openai(messages)
+        logging.info(f"Cleaned messages: {messages}")
 
         # Convert max_tokens to max_completion_tokens for newer models
         if "max_tokens" in kwargs:
@@ -182,6 +184,7 @@ class OpenAILLM(BaseLLM):
 
         try:
             for line in response:
+                logging.debug(f"OpenAI stream line: {line}")
                 if (
                     len(line.choices) > 0
                     and line.choices[0].delta.content is not None