Fix review comments and formatting

app.py

@@ -16,52 +16,11 @@ import ollama
 import requests
 import streamlit as st
 from dotenv import load_dotenv
-# Custom message classes to replace LangChain components
-class ChatMessage:
-    def __init__(self, content: str, role: str):
-        self.content = content
-        self.role = role
-        self.type = role  # For compatibility with existing code
-
-class HumanMessage(ChatMessage):
-    def __init__(self, content: str):
-        super().__init__(content, "user")
-
-class AIMessage(ChatMessage):
-    def __init__(self, content: str):
-        super().__init__(content, "ai")
-
-class StreamlitChatMessageHistory:
-    def __init__(self, key: str):
-        self.key = key
-        if key not in st.session_state:
-            st.session_state[key] = []
-    
-    @property
-    def messages(self):
-        return st.session_state[self.key]
-    
-    def add_user_message(self, content: str):
-        st.session_state[self.key].append(HumanMessage(content))
-    
-    def add_ai_message(self, content: str):
-        st.session_state[self.key].append(AIMessage(content))
-
-class ChatPromptTemplate:
-    def __init__(self, template: str):
-        self.template = template
-    
-    @classmethod
-    def from_template(cls, template: str):
-        return cls(template)
-    
-    def format(self, **kwargs):
-        return self.template.format(**kwargs)
 
 import global_config as gcfg
 import helpers.file_manager as filem
 from global_config import GlobalConfig
-from helpers import llm_helper, pptx_helper, text_helper
+from helpers import chat_helper, llm_helper, pptx_helper, text_helper
 
 load_dotenv()
 
@@ -333,8 +292,8 @@ def set_up_chat_ui():
     st.info(APP_TEXT['like_feedback'])
     st.chat_message('ai').write(random.choice(APP_TEXT['ai_greetings']))
 
-    history = StreamlitChatMessageHistory(key=CHAT_MESSAGES)
-    prompt_template = ChatPromptTemplate.from_template(
+    history = chat_helper.StreamlitChatMessageHistory(key=CHAT_MESSAGES)
+    prompt_template = chat_helper.ChatPromptTemplate.from_template(
         _get_prompt_template(
             is_refinement=_is_it_refinement()
         )
@@ -653,7 +612,7 @@ def _get_user_messages() -> List[str]:
     """
 
     return [
-        msg.content for msg in st.session_state[CHAT_MESSAGES] if isinstance(msg, HumanMessage)
+        msg.content for msg in st.session_state[CHAT_MESSAGES] if isinstance(msg, chat_helper.HumanMessage)
     ]
 
 

helpers/chat_helper.py

@@ -0,0 +1,60 @@
+"""
+Chat helper classes to replace LangChain components.
+"""
+import streamlit as st
+
+
+class ChatMessage:
+    """Base class for chat messages."""
+    
+    def __init__(self, content: str, role: str):
+        self.content = content
+        self.role = role
+        self.type = role  # For compatibility with existing code
+
+
+class HumanMessage(ChatMessage):
+    """Message from human user."""
+    
+    def __init__(self, content: str):
+        super().__init__(content, 'user')
+
+
+class AIMessage(ChatMessage):
+    """Message from AI assistant."""
+    
+    def __init__(self, content: str):
+        super().__init__(content, 'ai')
+
+
+class StreamlitChatMessageHistory:
+    """Chat message history stored in Streamlit session state."""
+    
+    def __init__(self, key: str):
+        self.key = key
+        if key not in st.session_state:
+            st.session_state[key] = []
+    
+    @property
+    def messages(self):
+        return st.session_state[self.key]
+    
+    def add_user_message(self, content: str):
+        st.session_state[self.key].append(HumanMessage(content))
+    
+    def add_ai_message(self, content: str):
+        st.session_state[self.key].append(AIMessage(content))
+
+
+class ChatPromptTemplate:
+    """Template for chat prompts."""
+    
+    def __init__(self, template: str):
+        self.template = template
+    
+    @classmethod
+    def from_template(cls, template: str):
+        return cls(template)
+    
+    def format(self, **kwargs):
+        return self.template.format(**kwargs)

helpers/llm_helper.py

@@ -73,17 +73,23 @@ def get_provider_model(provider_model: str, use_ollama: bool) -> Tuple[str, str]
             
             # Validate that the provider is in the valid providers list
             if inside_brackets not in GlobalConfig.VALID_PROVIDERS:
-                logger.warning(f"Provider '{inside_brackets}' not in VALID_PROVIDERS: {GlobalConfig.VALID_PROVIDERS}")
+                logger.warning(
+                    "Provider '%s' not in VALID_PROVIDERS: %s",
+                    inside_brackets, GlobalConfig.VALID_PROVIDERS
+                )
                 return '', ''
             
             # Validate that the model name is not empty
             if not outside_brackets.strip():
-                logger.warning(f"Empty model name for provider '{inside_brackets}'")
+                logger.warning("Empty model name for provider '%s'", inside_brackets)
                 return '', ''
             
             return inside_brackets, outside_brackets
 
-    logger.warning(f"Could not parse provider_model: '{provider_model}' (use_ollama={use_ollama})")
+    logger.warning(
+        "Could not parse provider_model: '%s' (use_ollama=%s)",
+        provider_model, use_ollama
+    )
     return '', ''
 
 
@@ -135,38 +141,20 @@ def get_litellm_model_name(provider: str, model: str) -> str:
     Convert provider and model to LiteLLM model name format.
     """
     provider_prefix_map = {
-        GlobalConfig.PROVIDER_HUGGING_FACE: "huggingface",
-        GlobalConfig.PROVIDER_GOOGLE_GEMINI: "gemini",
-        GlobalConfig.PROVIDER_AZURE_OPENAI: "azure",
-        GlobalConfig.PROVIDER_OPENROUTER: "openrouter",
-        GlobalConfig.PROVIDER_COHERE: "cohere",
-        GlobalConfig.PROVIDER_TOGETHER_AI: "together_ai",
-        GlobalConfig.PROVIDER_OLLAMA: "ollama",
+        GlobalConfig.PROVIDER_HUGGING_FACE: 'huggingface',
+        GlobalConfig.PROVIDER_GOOGLE_GEMINI: 'gemini',
+        GlobalConfig.PROVIDER_AZURE_OPENAI: 'azure',
+        GlobalConfig.PROVIDER_OPENROUTER: 'openrouter',
+        GlobalConfig.PROVIDER_COHERE: 'cohere',
+        GlobalConfig.PROVIDER_TOGETHER_AI: 'together_ai',
+        GlobalConfig.PROVIDER_OLLAMA: 'ollama',
     }
     prefix = provider_prefix_map.get(provider)
     if prefix:
-        return f"{prefix}/{model}"
+        return '%s/%s' % (prefix, model)
     return model
 
 
-def get_litellm_api_key(provider: str, api_key: str) -> str:
-    """
-    Get the appropriate API key for LiteLLM based on provider.
-    """
-    # All listed providers just return the api_key, so we can use a set for clarity
-    providers_with_api_key = {
-        GlobalConfig.PROVIDER_OPENROUTER,
-        GlobalConfig.PROVIDER_COHERE,
-        GlobalConfig.PROVIDER_TOGETHER_AI,
-        GlobalConfig.PROVIDER_GOOGLE_GEMINI,
-        GlobalConfig.PROVIDER_AZURE_OPENAI,
-        GlobalConfig.PROVIDER_HUGGING_FACE,
-    }
-    if provider in providers_with_api_key:
-        return api_key
-    return api_key
-
-
 def stream_litellm_completion(
         provider: str,
         model: str,
@@ -200,34 +188,32 @@ def stream_litellm_completion(
         # This is consistent with Azure OpenAI's requirement to use deployment names
         if not azure_deployment_name:
             raise ValueError("Azure deployment name is required for Azure OpenAI provider")
-        litellm_model = f"azure/{azure_deployment_name}"
+        litellm_model = 'azure/%s' % azure_deployment_name
     else:
         litellm_model = get_litellm_model_name(provider, model)
     
     # Prepare the request parameters
     request_params = {
-        "model": litellm_model,
-        "messages": messages,
-        "max_tokens": max_tokens,
-        "temperature": GlobalConfig.LLM_MODEL_TEMPERATURE,
-        "stream": True,
+        'model': litellm_model,
+        'messages': messages,
+        'max_tokens': max_tokens,
+        'temperature': GlobalConfig.LLM_MODEL_TEMPERATURE,
+        'stream': True,
     }
     
     # Set API key and any provider-specific params
     if provider != GlobalConfig.PROVIDER_OLLAMA:
-        # For OpenRouter, set environment variable as per documentation
+        # For OpenRouter, pass API key as parameter
         if provider == GlobalConfig.PROVIDER_OPENROUTER:
-            os.environ["OPENROUTER_API_KEY"] = api_key
-            # Don't add API key to request_params for OpenRouter
+            request_params['api_key'] = api_key
         elif provider == GlobalConfig.PROVIDER_AZURE_OPENAI:
-            # For Azure OpenAI, set environment variables as per documentation
-            os.environ["AZURE_API_KEY"] = api_key
-            os.environ["AZURE_API_BASE"] = azure_endpoint_url
-            os.environ["AZURE_API_VERSION"] = azure_api_version
+            # For Azure OpenAI, pass credentials as parameters
+            request_params['api_key'] = api_key
+            request_params['azure_api_base'] = azure_endpoint_url
+            request_params['azure_api_version'] = azure_api_version
         else:
             # For other providers, pass API key as parameter
-            api_key_to_use = get_litellm_api_key(provider, api_key)
-            request_params["api_key"] = api_key_to_use
+            request_params['api_key'] = api_key
     
     logger.debug('Streaming completion via LiteLLM: %s', litellm_model)
     
@@ -245,7 +231,7 @@ def stream_litellm_completion(
                         yield choice.message.content
                         
     except Exception as e:
-        logger.error(f"Error in LiteLLM completion: {e}")
+        logger.error('Error in LiteLLM completion: %s', e)
         raise
 
 
@@ -277,7 +263,10 @@ def get_litellm_llm(
     
     # Create a simple wrapper object that mimics the LangChain streaming interface
     class LiteLLMWrapper:
-        def __init__(self, provider, model, max_tokens, api_key, azure_endpoint_url, azure_deployment_name, azure_api_version):
+        def __init__(
+                self, provider, model, max_tokens, api_key, azure_endpoint_url,
+                azure_deployment_name, azure_api_version
+        ):
             self.provider = provider
             self.model = model
             self.max_tokens = max_tokens
@@ -287,7 +276,7 @@ def get_litellm_llm(
             self.azure_api_version = azure_api_version
         
         def stream(self, prompt: str):
-            messages = [{"role": "user", "content": prompt}]
+            messages = [{'role': 'user', 'content': prompt}]
             return stream_litellm_completion(
                 provider=self.provider,
                 model=self.model,
@@ -323,4 +312,4 @@ if __name__ == '__main__':
     ]
 
     for text in inputs:
-        print(get_provider_model(text, use_ollama=False))
+        print(get_provider_model(text, use_ollama=False))