Update app.py

app.py

@@ -19,7 +19,7 @@ from llama_cpp_agent.tools import WebSearchTool
 from llama_cpp_agent.prompt_templates import web_search_system_prompt, research_system_prompt
 from langchain_community.llms import HuggingFaceHub
 from llama_cpp_agent.llm_output_settings import LlmStructuredOutputSettings, LlmStructuredOutputType
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from llama_cpp_agent.llm_output_settings import LlmStructuredOutputType
 
 print("Available LlmStructuredOutputType options:")
@@ -103,10 +103,11 @@ examples = [
 
 class CustomLLMSettings(BaseModel):
     structured_output: LlmStructuredOutputSettings
-    temperature: float
-    top_p: float
-    repetition_penalty: float
-    
+    temperature: float = Field(default=0.7)
+    top_p: float = Field(default=0.95)
+    repetition_penalty: float = Field(default=1.1)
+    top_k: int = Field(default=50)  # Added top_k parameter
+
 class HuggingFaceHubWrapper:
     def __init__(self, repo_id, model_kwargs, huggingfacehub_api_token):
         self.model = HuggingFaceHub(
@@ -117,7 +118,7 @@ class HuggingFaceHubWrapper:
         self.temperature = model_kwargs.get('temperature', 0.7)
         self.top_p = model_kwargs.get('top_p', 0.95)
         self.repetition_penalty = model_kwargs.get('repetition_penalty', 1.1)
-
+        self.top_k = model_kwargs.get('top_k', 50)  # Added top_k
 
     def get_provider_default_settings(self):
         return CustomLLMSettings(
@@ -129,7 +130,8 @@ class HuggingFaceHubWrapper:
             ),
             temperature=self.temperature,
             top_p=self.top_p,
-            repetition_penalty=self.repetition_penalty
+            repetition_penalty=self.repetition_penalty,
+            top_k=self.top_k  # Added top_k
         )
 
     def get_provider_identifier(self):
@@ -170,13 +172,14 @@ class CitingSources(BaseModel):
     )
 
 # Model function
-def get_model(temperature, top_p, repetition_penalty):
+def get_model(temperature, top_p, repetition_penalty, top_k=50):
     return HuggingFaceHubWrapper(
         repo_id="mistralai/Mistral-7B-Instruct-v0.3",
         model_kwargs={
             "temperature": temperature,
             "top_p": top_p,
             "repetition_penalty": repetition_penalty,
+            "top_k": top_k,
             "max_length": 1000
         },
         huggingfacehub_api_token=huggingface_token
@@ -204,8 +207,9 @@ def respond(
     temperature,
     top_p,
     repeat_penalty,
+    top_k=50,  # Added top_k parameter
 ):
-    model = get_model(temperature, top_p, repeat_penalty)
+    model = get_model(temperature, top_p, repeat_penalty, top_k)
     
     chat_template = MessagesFormatterType.MISTRAL
     
@@ -258,6 +262,7 @@ demo = gr.ChatInterface(
         gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
         gr.Slider(minimum=0.0, maximum=2.0, value=1.1, step=0.1, label="Repetition penalty"),
+        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k"),  # Added top_k slider
     ],
     theme=gr.themes.Soft(
         primary_hue="orange",