Update app.py

app.py

@@ -13,7 +13,7 @@ import base64
 import logging
 import time
 from urllib.parse import quote  # Added for URL encoding
-import importlib  # NEW: For dynamic import
+import importlib  # For dynamic import
 
 import gradio as gr
 import spaces
@@ -84,7 +84,6 @@ def generate_image(prompt: str, width: float, height: float, guidance: float, in
         logging.error(f"Image generation failed: {str(e)}")
         return None, f"Error: {str(e)}"
 
-# Base64 padding fix function
 def fix_base64_padding(data):
     """Fix the padding of a Base64 string."""
     if isinstance(data, bytes):
@@ -99,18 +98,12 @@ def fix_base64_padding(data):
     
     return data
 
-# =============================================================================
-# Memory cleanup function
-# =============================================================================
 def clear_cuda_cache():
     """Explicitly clear the CUDA cache."""
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         gc.collect()
 
-# =============================================================================
-# SerpHouse related functions
-# =============================================================================
 SERPHOUSE_API_KEY = os.getenv("SERPHOUSE_API_KEY", "")
 
 def extract_keywords(text: str, top_k: int = 5) -> str:
@@ -176,9 +169,6 @@ Below are the search results. Use this information to answer the query:
         logger.error(f"Web search failed: {e}")
         return f"Web search failed: {str(e)}"
 
-# =============================================================================
-# Model and processor loading
-# =============================================================================
 MAX_CONTENT_CHARS = 2000
 MAX_INPUT_LENGTH = 2096
 model_id = os.getenv("MODEL_ID", "VIDraft/Gemma-3-R1984-4B")
@@ -191,9 +181,6 @@ model = Gemma3ForConditionalGeneration.from_pretrained(
 )
 MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))
 
-# =============================================================================
-# CSV, TXT, PDF analysis functions
-# =============================================================================
 def analyze_csv_file(path: str) -> str:
     try:
         df = pd.read_csv(path)
@@ -238,9 +225,6 @@ def pdf_to_markdown(pdf_path: str) -> str:
         full_text = full_text[:MAX_CONTENT_CHARS] + "\n...(truncated)..."
     return f"**[PDF File: {os.path.basename(pdf_path)}]**\n\n{full_text}"
 
-# =============================================================================
-# Check media file limits
-# =============================================================================
 def count_files_in_new_message(paths: list[str]) -> tuple[int, int]:
     image_count = 0
     video_count = 0
@@ -293,9 +277,6 @@ def validate_media_constraints(message: dict, history: list[dict]) -> bool:
             return False
     return True
 
-# =============================================================================
-# Video processing functions
-# =============================================================================
 def downsample_video(video_path: str) -> list[tuple[Image.Image, float]]:
     vidcap = cv2.VideoCapture(video_path)
     fps = vidcap.get(cv2.CAP_PROP_FPS)
@@ -328,9 +309,6 @@ def process_video(video_path: str) -> tuple[list[dict], list[str]]:
             content.append({"type": "image", "url": temp_file.name})
     return content, temp_files
 
-# =============================================================================
-# Interleaved <image> processing function
-# =============================================================================
 def process_interleaved_images(message: dict) -> list[dict]:
     parts = re.split(r"(<image>)", message["text"])
     content = []
@@ -347,9 +325,6 @@ def process_interleaved_images(message: dict) -> list[dict]:
                 content.append({"type": "text", "text": part})
     return content
 
-# =============================================================================
-# File processing -> content creation
-# =============================================================================
 def is_image_file(file_path: str) -> bool:
     return bool(re.search(r"\.(png|jpg|jpeg|gif|webp)$", file_path, re.IGNORECASE))
 
@@ -390,9 +365,6 @@ def process_new_user_message(message: dict) -> tuple[list[dict], list[str]]:
             content_list.append({"type": "image", "url": img_path})
     return content_list, temp_files
 
-# =============================================================================
-# Convert history to LLM messages
-# =============================================================================
 def process_history(history: list[dict]) -> list[dict]:
     messages = []
     current_user_content = []
@@ -416,9 +388,6 @@ def process_history(history: list[dict]) -> list[dict]:
         messages.append({"role": "user", "content": current_user_content})
     return messages
 
-# =============================================================================
-# Model generation function (with OOM catching)
-# =============================================================================
 def _model_gen_with_oom_catch(**kwargs):
     try:
         model.generate(**kwargs)
@@ -433,18 +402,10 @@ def _model_gen_with_oom_catch(**kwargs):
 def load_function_definitions(json_path="functions.json"):
     """
     로컬 JSON 파일에서 함수 정의 목록을 로드하여 반환.
-    각 항목: {
-      "name": <str>,
-      "description": <str>,
-      "module_path": <str>,
-      "func_name_in_module": <str>,
-      "parameters": { ... }
-    }
     """
     try:
         with open(json_path, "r", encoding="utf-8") as f:
             data = json.load(f)
-        # name을 키로 하는 dict 형태로 재구성
         func_dict = {}
         for entry in data:
             func_name = entry["name"]
@@ -456,9 +417,6 @@ def load_function_definitions(json_path="functions.json"):
 
 FUNCTION_DEFINITIONS = load_function_definitions("functions.json")
 
-# =============================================================================
-# Dynamic handle_function_call
-# =============================================================================
 def handle_function_call(text: str) -> str:
     """
     Detects and processes function call blocks in the text using the JSON-based approach.
@@ -470,7 +428,6 @@ def handle_function_call(text: str) -> str:
     ```tool_code
     get_product_name_by_PID(PID="807ZPKBL9V")
     ```
-    We parse that block, check if 'FUNCTION_DEFINITIONS' has an entry, then import & call it.
     """
     import re
     pattern = r"```tool_code\s*(.*?)\s*```"
@@ -479,12 +436,11 @@ def handle_function_call(text: str) -> str:
         return ""
     code_block = match.group(1).strip()
 
-    # 함수명 추출 (예: get_stock_price)
-    # 정규식: ^(\w+)\(.*\)
     func_match = re.match(r'^(\w+)\((.*)\)$', code_block)
     if not func_match:
         logger.debug("No valid function call format found.")
         return ""
+
     func_name = func_match.group(1)
     param_str = func_match.group(2).strip()
 
@@ -496,43 +452,35 @@ def handle_function_call(text: str) -> str:
     func_info = FUNCTION_DEFINITIONS[func_name]
     module_path = func_info["module_path"]
     module_func_name = func_info["func_name_in_module"]
-    # 동적 임포트
+
     try:
         imported_module = importlib.import_module(module_path)
     except ImportError as e:
         logger.error(f"Failed to import module {module_path}: {e}")
         return f"```tool_output\nError: Cannot import module '{module_path}'\n```"
 
-    # 실제 함수 객체를 가져옴
     if not hasattr(imported_module, module_func_name):
         logger.error(f"Module '{module_path}' has no attribute '{module_func_name}'.")
         return f"```tool_output\nError: Function '{module_func_name}' not found in module '{module_path}'\n```"
 
     real_func = getattr(imported_module, module_func_name)
 
-    # 파라미터 파싱 예: ticker="AAPL", some_arg=123
-    # 단순 정규식으로 key="value" or key=123 식을 구분
+    # 간단 파라미터 파싱 (key="value" or key=123)
     param_pattern = r'(\w+)\s*=\s*"(.*?)"|(\w+)\s*=\s*([\d.]+)'
-    # 이 정규식은 간단히 key="string" 또는 key=123 같은 형태를 파싱
-    # 더 복잡한 경우 별도 파싱 로직이나 json.loads 기법 사용 필요
     param_dict = {}
     for p_match in re.finditer(param_pattern, param_str):
         if p_match.group(1) and p_match.group(2):
-            # group(1)은 key, group(2)는 string value
             key = p_match.group(1)
             val = p_match.group(2)
             param_dict[key] = val
         else:
-            # group(3)은 key, group(4)는 numeric value
             key = p_match.group(3)
             val = p_match.group(4)
-            # 숫자 변환
             if '.' in val:
                 param_dict[key] = float(val)
             else:
                 param_dict[key] = int(val)
 
-    # 이제 실제 함수 실행
     try:
         result = real_func(**param_dict)
     except Exception as e:
@@ -541,9 +489,6 @@ def handle_function_call(text: str) -> str:
 
     return f"```tool_output\n{result}\n```"
 
-# =============================================================================
-# Main inference function
-# =============================================================================
 @spaces.GPU(duration=120)
 def run(
     message: dict,
@@ -555,19 +500,18 @@ def run(
     age_group: str = "20s",
     mbti_personality: str = "INTP",
     sexual_openness: int = 2,
-    image_gen: bool = False  # "Image Gen" checkbox status
+    image_gen: bool = False
 ) -> Iterator[str]:
     if not validate_media_constraints(message, history):
         yield ""
         return
     temp_files = []
     try:
-        # JSON에서 로드된 함수 목록을 요약해서 시스템 프롬프트에 포함할 수도 있음
-        # (토큰 부담이 커질 수 있으므로, 적당히 압축 요약 권장)
-        # 아래는 예시로 간단히 함수 이름만 나열
+        # JSON에서 로드된 함수 정보 문자열화 (예: 함수명과 example_usage만)
         available_funcs_text = ""
         for f_name, info in FUNCTION_DEFINITIONS.items():
-            available_funcs_text += f"Function: {f_name} - {info['description']}\n"
+            example_usage = info.get("example_usage", "")
+            available_funcs_text += f"\n\nFunction: {f_name}\nDescription: {info['description']}\nExample:\n{example_usage}\n"
 
         persona = (
             f"{system_prompt.strip()}\n\n"
@@ -575,7 +519,9 @@ def run(
             f"Age Group: {age_group}\n"
             f"MBTI Persona: {mbti_personality}\n"
             f"Sexual Openness (1-5): {sexual_openness}\n\n"
-            "Below are the available functions you can call (use the format: ```tool_code\\nfunc_name(param=...)\n```):\n"
+            "Below are the available functions you can call.\n"
+            "Important: Use the format exactly like: ```tool_code\nfunctionName(param=\"string\", ...)\n```\n"
+            "(Strings must be in double quotes)\n"
             f"{available_funcs_text}\n"
         )
         combined_system_msg = f"[System Prompt]\n{persona.strip()}\n\n"
@@ -629,7 +575,6 @@ def run(
             output_so_far += new_text
             yield output_so_far
 
-        # 모델 출력 중 ```tool_code``` 블록이 있으면 처리
         func_result = handle_function_call(output_so_far)
         if func_result:
             output_so_far += "\n\n" + func_result
@@ -652,17 +597,12 @@ def run(
             pass
         clear_cuda_cache()
 
-# =============================================================================
-# Modified model run function - handles image generation and gallery update
-# =============================================================================
 def modified_run(message, history, system_prompt, max_new_tokens, use_web_search, web_search_query,
                 age_group, mbti_personality, sexual_openness, image_gen):
-    # Initialize and hide the gallery component
     output_so_far = ""
     gallery_update = gr.Gallery(visible=False, value=[])
     yield output_so_far, gallery_update
     
-    # Execute the original run function
     text_generator = run(message, history, system_prompt, max_new_tokens, use_web_search, 
                        web_search_query, age_group, mbti_personality, sexual_openness, image_gen)
     
@@ -670,15 +610,12 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
         output_so_far = text_chunk
         yield output_so_far, gallery_update
         
-    # If image generation is enabled and there is text input, update the gallery
     if image_gen and message["text"].strip():
         try:
             width, height = 512, 512
             guidance, steps, seed = 7.5, 30, 42
             
             logger.info(f"Calling image generation for gallery with prompt: {message['text']}")
-            
-            # Call the API to generate an image
             image_result, seed_info = generate_image(
                 prompt=message["text"].strip(),
                 width=width, 
@@ -687,7 +624,6 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
                 inference_steps=steps, 
                 seed=seed
             )
-            
             if image_result:
                 if isinstance(image_result, str) and (
                     image_result.startswith('data:') or 
@@ -699,22 +635,18 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
                         else:
                             b64data = image_result
                             content_type = "image/webp"
-
                         image_bytes = base64.b64decode(b64data)
                         with tempfile.NamedTemporaryFile(delete=False, suffix=".webp") as temp_file:
                             temp_file.write(image_bytes)
                             temp_path = temp_file.name
                             gallery_update = gr.Gallery(visible=True, value=[temp_path])
                             yield output_so_far + "\n\n*Image generated and displayed in the gallery below.*", gallery_update
-                            
                     except Exception as e:
                         logger.error(f"Error processing Base64 image: {e}")
                         yield output_so_far + f"\n\n(Error processing image: {e})", gallery_update
-                        
                 elif isinstance(image_result, str) and os.path.exists(image_result):
                     gallery_update = gr.Gallery(visible=True, value=[image_result])
                     yield output_so_far + "\n\n*Image generated and displayed in the gallery below.*", gallery_update
-                
                 elif isinstance(image_result, str) and '/tmp/' in image_result:
                     try:
                         client = Client(API_URL)
@@ -722,13 +654,11 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
                             prompt=message["text"].strip(),
                             api_name="/generate_base64_image"
                         )
-                        
                         if isinstance(result, str) and (result.startswith('data:') or len(result) > 100):
                             if result.startswith('data:'):
                                 content_type, b64data = result.split(';base64,')
                             else:
                                 b64data = result
-                            
                             image_bytes = base64.b64decode(b64data)
                             with tempfile.NamedTemporaryFile(delete=False, suffix=".webp") as temp_file:
                                 temp_file.write(image_bytes)
@@ -737,7 +667,6 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
                                 yield output_so_far + "\n\n*Image generated and displayed in the gallery below.*", gallery_update
                         else:
                             yield output_so_far + "\n\n(Image generation failed: Invalid format)", gallery_update
-                            
                     except Exception as e:
                         logger.error(f"Error calling alternative API: {e}")
                         yield output_so_far + f"\n\n(Image generation failed: {e})", gallery_update
@@ -755,14 +684,10 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
                     yield output_so_far + f"\n\n(Unsupported image format: {type(image_result)})", gallery_update
             else:
                 yield output_so_far + f"\n\n(Image generation failed: {seed_info})", gallery_update
-                
         except Exception as e:
             logger.error(f"Error during gallery image generation: {e}")
             yield output_so_far + f"\n\n(Image generation error: {e})", gallery_update
 
-# =============================================================================
-# Examples
-# =============================================================================
 examples = [
     [
         {
@@ -855,7 +780,7 @@ examples = [
     ],
     [
         {
-            "text": "AAPL의 현재 주가를 알려줘.",  
+            "text": "AAPL의 현재 주가를 알려줘.",
             "files": []
         }
     ],