Update app/webui/patch.py

app/webui/patch.py

@@ -1,131 +1,131 @@
-# a monkey patch to use llama-index completion
-from typing import Union, Callable
-from functools import wraps
-from src.translation_agent.utils import *
-
-
-from llama_index.llms.groq import Groq
-from llama_index.llms.cohere import Cohere
-from llama_index.llms.openai import OpenAI
-from llama_index.llms.together import TogetherLLM
-from llama_index.llms.ollama import Ollama
-from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
-
-from llama_index.core import Settings
-from llama_index.core.llms import ChatMessage
-
-
-# Add your LLMs here
-
-def model_load(
-        endpoint: str,
-        model: str,
-        api_key: str = None,
-        context_window: int = 4096,
-        num_output: int = 512,
-):
-    if endpoint == "Groq":
-        llm = Groq(
-            model=model,
-            api_key=api_key,
-        )
-    elif endpoint == "Cohere":
-        llm = Cohere(
-            model=model,
-            api_key=api_key,
-        )
-    elif endpoint == "OpenAI":
-        llm = OpenAI(
-            model=model,
-            api_key=api_key,
-        )
-    elif endpoint == "TogetherAI":
-        llm = TogetherLLM(
-            model=model,
-            api_key=api_key,
-        )
-    elif endpoint == "ollama":
-        llm = Ollama(
-            model=model,
-            request_timeout=120.0)
-    elif endpoint == "Huggingface":
-        llm = HuggingFaceInferenceAPI(
-            model_name=model,
-            token=api_key,
-            task="text-generation",
-        )
-    Settings.llm = llm
-    # maximum input size to the LLM
-    Settings.context_window = context_window
-
-    # number of tokens reserved for text generation.
-    Settings.num_output = num_output
-
-
-
-def completion_wrapper(func: Callable) -> Callable:
-    @wraps(func)
-    def wrapper(
-        prompt: str,
-        system_message: str = "You are a helpful assistant.",
-        temperature: float = 0.3,
-        json_mode: bool = False,
-    ) -> Union[str, dict]:
-        """
-            Generate a completion using the OpenAI API.
-
-        Args:
-            prompt (str): The user's prompt or query.
-            system_message (str, optional): The system message to set the context for the assistant.
-                Defaults to "You are a helpful assistant.".
-            temperature (float, optional): The sampling temperature for controlling the randomness of the generated text.
-                Defaults to 0.3.
-            json_mode (bool, optional): Whether to return the response in JSON format.
-                Defaults to False.
-
-        Returns:
-            Union[str, dict]: The generated completion.
-                If json_mode is True, returns the complete API response as a dictionary.
-                If json_mode is False, returns the generated text as a string.
-        """
-        llm = Settings.llm
-        if llm.class_name() == "HuggingFaceInferenceAPI":
-            llm.system_prompt = system_message
-            messages = [
-                ChatMessage(
-                    role="user", content=prompt),
-            ]
-            response = llm.chat(
-                messages=messages,
-                temperature=temperature,
-                top_p=1,
-            )
-            return response.message.content
-        else:
-            messages = [
-                ChatMessage(
-                    role="system", content=system_message),
-                ChatMessage(
-                    role="user", content=prompt),
-            ]
-
-            if json_mode:
-                response = llm.chat(
-                    temperature=temperature,
-                    top_p=1,
-                    response_format={"type": "json_object"},
-                    messages=messages,
-                )
-                return response.message.content
-            else:
-                response = llm.chat(
-                    temperature=temperature,
-                    top_p=1,
-                    messages=messages,
-                )
-                return response.message.content
-
-    return wrapper
-
-openai_completion = get_completion
+# a monkey patch to use llama-index completion
+from typing import Union, Callable
+from functools import wraps
+from src.translation_agent.utils import *
+
+
+from llama_index.llms.groq import Groq
+from llama_index.llms.cohere import Cohere
+from llama_index.llms.openai import OpenAI
+from llama_index.llms.together import TogetherLLM
+from llama_index.llms.ollama import Ollama
+from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
+
+from llama_index.core import Settings
+from llama_index.core.llms import ChatMessage
+
+
+# Add your LLMs here
+
+def model_load(
+        endpoint: str,
+        model: str,
+        api_key: str = None,
+        context_window: int = 4096,
+        num_output: int = 512,
+):
+    if endpoint == "Groq":
+        llm = Groq(
+            model=model,
+            api_key=api_key,
+        )
+    elif endpoint == "Cohere":
+        llm = Cohere(
+            model=model,
+            api_key=api_key,
+        )
+    elif endpoint == "OpenAI":
+        llm = OpenAI(
+            model=model,
+            api_key=api_key if api_key else os.getenv("OPENAI_API_KEY"),
+        )
+    elif endpoint == "TogetherAI":
+        llm = TogetherLLM(
+            model=model,
+            api_key=api_key,
+        )
+    elif endpoint == "ollama":
+        llm = Ollama(
+            model=model,
+            request_timeout=120.0)
+    elif endpoint == "Huggingface":
+        llm = HuggingFaceInferenceAPI(
+            model_name=model,
+            token=api_key,
+            task="text-generation",
+        )
+    Settings.llm = llm
+    # maximum input size to the LLM
+    Settings.context_window = context_window
+
+    # number of tokens reserved for text generation.
+    Settings.num_output = num_output
+
+
+
+def completion_wrapper(func: Callable) -> Callable:
+    @wraps(func)
+    def wrapper(
+        prompt: str,
+        system_message: str = "You are a helpful assistant.",
+        temperature: float = 0.3,
+        json_mode: bool = False,
+    ) -> Union[str, dict]:
+        """
+            Generate a completion using the OpenAI API.
+
+        Args:
+            prompt (str): The user's prompt or query.
+            system_message (str, optional): The system message to set the context for the assistant.
+                Defaults to "You are a helpful assistant.".
+            temperature (float, optional): The sampling temperature for controlling the randomness of the generated text.
+                Defaults to 0.3.
+            json_mode (bool, optional): Whether to return the response in JSON format.
+                Defaults to False.
+
+        Returns:
+            Union[str, dict]: The generated completion.
+                If json_mode is True, returns the complete API response as a dictionary.
+                If json_mode is False, returns the generated text as a string.
+        """
+        llm = Settings.llm
+        if llm.class_name() == "HuggingFaceInferenceAPI":
+            llm.system_prompt = system_message
+            messages = [
+                ChatMessage(
+                    role="user", content=prompt),
+            ]
+            response = llm.chat(
+                messages=messages,
+                temperature=temperature,
+                top_p=1,
+            )
+            return response.message.content
+        else:
+            messages = [
+                ChatMessage(
+                    role="system", content=system_message),
+                ChatMessage(
+                    role="user", content=prompt),
+            ]
+
+            if json_mode:
+                response = llm.chat(
+                    temperature=temperature,
+                    top_p=1,
+                    response_format={"type": "json_object"},
+                    messages=messages,
+                )
+                return response.message.content
+            else:
+                response = llm.chat(
+                    temperature=temperature,
+                    top_p=1,
+                    messages=messages,
+                )
+                return response.message.content
+
+    return wrapper
+
+openai_completion = get_completion
 get_completion = completion_wrapper(openai_completion)