Update app.py

app.py

@@ -157,16 +157,58 @@ class _RC:
         _prompt += "Assistant: "
         return _prompt, _system
 
+    def _sanitize_params(self, **_kwargs):
+        """Sanitize parameters and set proper defaults"""
+        _params = {}
+
+        # Handle max_tokens
+        _max_tokens = _kwargs.get('max_tokens')
+        if _max_tokens is not None and _max_tokens > 0:
+            _params['max_tokens'] = _max_tokens
+        else:
+            _params['max_tokens'] = 4096
+
+        # Handle temperature
+        _temperature = _kwargs.get('temperature')
+        if _temperature is not None:
+            _params['temperature'] = max(0.0, min(2.0, float(_temperature)))
+        else:
+            _params['temperature'] = 0.7
+
+        # Handle top_p
+        _top_p = _kwargs.get('top_p')
+        if _top_p is not None:
+            _params['top_p'] = max(0.0, min(1.0, float(_top_p)))
+        else:
+            _params['top_p'] = 1.0
+
+        # Handle presence_penalty
+        _presence_penalty = _kwargs.get('presence_penalty')
+        if _presence_penalty is not None:
+            _params['presence_penalty'] = max(-2.0, min(2.0, float(_presence_penalty)))
+        else:
+            _params['presence_penalty'] = 0.0
+
+        # Handle frequency_penalty
+        _frequency_penalty = _kwargs.get('frequency_penalty')
+        if _frequency_penalty is not None:
+            _params['frequency_penalty'] = max(-2.0, min(2.0, float(_frequency_penalty)))
+        else:
+            _params['frequency_penalty'] = 0.0
+
+        return _params
+
     def _create_prediction(self, _model_name, _prompt, _system="", **_kwargs):
         """Create a prediction using Replicate API"""
         _replicate_model = self._get_replicate_model(_model_name)
+        _params = self._sanitize_params(**_kwargs)
 
         _input = {
             "prompt": _prompt,
             "system_prompt": _system,
-            "max_tokens": _kwargs.get('max_tokens', 4096),
-            "temperature": _kwargs.get('temperature', 0.7),
-            "top_p": _kwargs.get('top_p', 1.0)
+            "max_tokens": _params['max_tokens'],
+            "temperature": _params['temperature'],
+            "top_p": _params['top_p']
         }
 
         try:
@@ -200,13 +242,14 @@ class _RC:
     def _stream_chat(self, _model_name, _prompt, _system="", **_kwargs):
         """Stream chat using Replicate's streaming API"""
         _replicate_model = self._get_replicate_model(_model_name)
+        _params = self._sanitize_params(**_kwargs)
 
         _input = {
             "prompt": _prompt,
             "system_prompt": _system,
-            "max_tokens": _kwargs.get('max_tokens', 4096),
-            "temperature": _kwargs.get('temperature', 0.7),
-            "top_p": _kwargs.get('top_p', 1.0)
+            "max_tokens": _params['max_tokens'],
+            "temperature": _params['temperature'],
+            "top_p": _params['top_p']
         }
 
         try:
@@ -253,13 +296,14 @@ class _RC:
     def _complete_chat(self, _model_name, _prompt, _system="", **_kwargs):
         """Complete chat using Replicate's run method"""
         _replicate_model = self._get_replicate_model(_model_name)
+        _params = self._sanitize_params(**_kwargs)
 
         _input = {
             "prompt": _prompt,
             "system_prompt": _system,
-            "max_tokens": _kwargs.get('max_tokens', 4096),
-            "temperature": _kwargs.get('temperature', 0.7),
-            "top_p": _kwargs.get('top_p', 1.0)
+            "max_tokens": _params['max_tokens'],
+            "temperature": _params['temperature'],
+            "top_p": _params['top_p']
         }
 
         try:
@@ -433,8 +477,17 @@ async def _generate_stream_response(_request: _CCR, _prompt: str, _system: str,
         _total_content = ""
 
         try:
+            # Extract only relevant parameters for Replicate API
+            _api_params = {
+                'max_tokens': _request.max_tokens,
+                'temperature': _request.temperature,
+                'top_p': _request.top_p,
+                'presence_penalty': _request.presence_penalty,
+                'frequency_penalty': _request.frequency_penalty
+            }
+
             # Use Replicate's direct streaming method with model parameter
-            for _chunk in _client._stream_chat(_request.model, _prompt, _system, **_request.model_dump()):
+            for _chunk in _client._stream_chat(_request.model, _prompt, _system, **_api_params):
                 if _chunk and isinstance(_chunk, str):
                     _chunk_count += 1
                     _total_content += _chunk
@@ -539,6 +592,17 @@ async def _create_chat_completion(_request: _CCR):
 
         _lg.info(f"[{_request_id}] Formatted prompt length: {len(_prompt)}")
 
+        # Extract only relevant parameters for Replicate API
+        _api_params = {
+            'max_tokens': _request.max_tokens,
+            'temperature': _request.temperature,
+            'top_p': _request.top_p,
+            'presence_penalty': _request.presence_penalty,
+            'frequency_penalty': _request.frequency_penalty
+        }
+
+        _lg.info(f"[{_request_id}] API parameters: {_api_params}")
+
         # Stream or complete
         if _request.stream:
             _lg.info(f"[{_request_id}] Starting streaming response")
@@ -554,7 +618,7 @@ async def _create_chat_completion(_request: _CCR):
         else:
             # Non-streaming completion
             _lg.info(f"[{_request_id}] Starting non-streaming completion")
-            _content = _client._complete_chat(_request.model, _prompt, _system, **_request.model_dump())
+            _content = _client._complete_chat(_request.model, _prompt, _system, **_api_params)
 
             _completion_id = f"chatcmpl-{_u.uuid4().hex}"
             _created_time = int(_t.time())