Update cohereAPI.py

cohereAPI.py

@@ -11,7 +11,7 @@ def get_client(api_key):
         _client = cohere.ClientV2(api_key)
     return _client
 
-def send_message_stream(system_message, user_message, conversation_history, api_key, model_name="command-a-03-2025"):
+def send_message_stream(system_message, user_message, conversation_history, api_key, model_name="command-a-03-2025", temperature=0.7, max_tokens=None):
     """Stream response from Cohere API"""
     # Get or create the Cohere client
     co = get_client(api_key)
@@ -21,11 +21,19 @@ def send_message_stream(system_message, user_message, conversation_history, api_
     messages.extend(conversation_history)
     messages.append({"role": "user", "content": user_message})
     
+    # Prepare chat parameters
+    chat_params = {
+        "model": model_name,
+        "messages": messages,
+        "temperature": temperature
+    }
+    
+    # Add max_tokens if specified
+    if max_tokens:
+        chat_params["max_tokens"] = int(max_tokens)
+    
     # Send streaming request to Cohere
-    stream = co.chat_stream(
-        model=model_name,
-        messages=messages
-    )
+    stream = co.chat_stream(**chat_params)
     
     # Collect full response for history
     full_response = ""
@@ -37,10 +45,10 @@ def send_message_stream(system_message, user_message, conversation_history, api_
             full_response += text_chunk
             yield text_chunk
 
-async def send_message_stream_async(system_message, user_message, conversation_history, api_key, model_name="command-a-03-2025"):
+async def send_message_stream_async(system_message, user_message, conversation_history, api_key, model_name="command-a-03-2025", temperature=0.7, max_tokens=None):
     """Async wrapper for streaming response from Cohere API"""
     def _sync_stream():
-        return send_message_stream(system_message, user_message, conversation_history, api_key, model_name)
+        return send_message_stream(system_message, user_message, conversation_history, api_key, model_name, temperature, max_tokens)
     
     # Run the synchronous generator in a thread
     loop = asyncio.get_event_loop()
@@ -72,7 +80,7 @@ async def send_message_stream_async(system_message, user_message, conversation_h
         yield chunk
     
 
-def send_message(system_message, user_message, conversation_history, api_key, model_name="command-a-03-2025"):
+def send_message(system_message, user_message, conversation_history, api_key, model_name="command-a-03-2025", temperature=0.7, max_tokens=None):
     """Non-streaming version for backward compatibility"""
     # Get or create the Cohere client
     co = get_client(api_key)
@@ -82,18 +90,26 @@ def send_message(system_message, user_message, conversation_history, api_key, mo
     messages.extend(conversation_history)
     messages.append({"role": "user", "content": user_message})
     
+    # Prepare chat parameters
+    chat_params = {
+        "model": model_name,
+        "messages": messages,
+        "temperature": temperature
+    }
+    
+    # Add max_tokens if specified
+    if max_tokens:
+        chat_params["max_tokens"] = int(max_tokens)
+    
     # Send request to Cohere synchronously
-    response = co.chat(
-        model=model_name,
-        messages=messages
-    )
+    response = co.chat(**chat_params)
     
     # Get the response
     response_content = response.message.content[0].text
     
     return response_content
 
-async def send_message_async(system_message, user_message, conversation_history, api_key, model_name="command-a-03-2025"):
+async def send_message_async(system_message, user_message, conversation_history, api_key, model_name="command-a-03-2025", temperature=0.7, max_tokens=None):
     """Async version using asyncio.to_thread"""
     return await asyncio.to_thread(
         send_message, 
@@ -101,5 +117,7 @@ async def send_message_async(system_message, user_message, conversation_history,
         user_message, 
         conversation_history, 
         api_key, 
-        model_name
+        model_name,
+        temperature,
+        max_tokens
     )