Update app.py

app.py

@@ -13,6 +13,7 @@ import base64
 import logging
 import time
 from urllib.parse import quote  # Added for URL encoding
+import importlib  # NEW: For dynamic import
 
 import gradio as gr
 import spaces
@@ -89,11 +90,9 @@ def fix_base64_padding(data):
     if isinstance(data, bytes):
         data = data.decode('utf-8')
     
-    # Remove the prefix if present
     if "base64," in data:
         data = data.split("base64,", 1)[1]
     
-    # Add padding characters (to make the length a multiple of 4)
     missing_padding = len(data) % 4
     if missing_padding:
         data += '=' * (4 - missing_padding)
@@ -429,61 +428,118 @@ def _model_gen_with_oom_catch(**kwargs):
         clear_cuda_cache()
 
 # =============================================================================
-# Yahoo Finance 함수: yfinance를 활용하여 주식 가격 조회
+# JSON 기반 함수 목록 로드
 # =============================================================================
-import yfinance as yf
-
-def get_stock_price(ticker: str) -> float:
+def load_function_definitions(json_path="functions.json"):
     """
-    주어진 티커(ticker)의 최신 종가를 반환합니다.
-    yfinance 라이브러리를 사용하며, 별도의 토큰 없이 데이터를 가져옵니다.
+    로컬 JSON 파일에서 함수 정의 목록을 로드하여 반환.
+    각 항목: {
+      "name": <str>,
+      "description": <str>,
+      "module_path": <str>,
+      "func_name_in_module": <str>,
+      "parameters": { ... }
+    }
     """
-    stock = yf.Ticker(ticker)
-    data = stock.history(period="1d")
-    if not data.empty:
-        return data['Close'].iloc[-1]
-    return float('nan')
+    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"]
+            func_dict[func_name] = entry
+        return func_dict
+    except Exception as e:
+        logger.error(f"Failed to load function definitions from JSON: {e}")
+        return {}
+
+FUNCTION_DEFINITIONS = load_function_definitions("functions.json")
 
 # =============================================================================
-# 함수 호출 예제: 제품 조회 및 주식 가격 조회 함수 처리
+# Dynamic handle_function_call
 # =============================================================================
-def get_product_name_by_PID(PID: str) -> str:
-    """Finds the name of a product by its Product ID"""
-    product_catalog = {
-        "807ZPKBL9V": "SuperWidget",
-        "1234567890": "MegaGadget"
-    }
-    return product_catalog.get(PID, "Unknown product")
-
 def handle_function_call(text: str) -> str:
     """
-    Detects and processes function call blocks in the text.
-    처리 대상:
-      - get_product_name_by_PID(PID="...")
-      - get_stock_price(ticker="...")
-    그리고 결과를 tool_output 블록으로 반환합니다.
+    Detects and processes function call blocks in the text using the JSON-based approach.
+    The model is expected to produce something like:
+    ```tool_code
+    get_stock_price(ticker="AAPL")
+    ```
+    or
+    ```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, io
-    from contextlib import redirect_stdout
+    import re
     pattern = r"```tool_code\s*(.*?)\s*```"
     match = re.search(pattern, text, re.DOTALL)
-    if match:
-        code = match.group(1).strip()
-        # 제품 조회 함수 처리
-        if code.startswith("get_product_name_by_PID("):
-            pid_match = re.search(r'PID\s*=\s*"(.*?)"', code)
-            if pid_match:
-                pid = pid_match.group(1)
-                result = get_product_name_by_PID(pid)
-                return f"```tool_output\n{result}\n```"
-        # 주식 가격 조회 함수 처리
-        elif code.startswith("get_stock_price("):
-            ticker_match = re.search(r'ticker\s*=\s*"(.*?)"', code)
-            if ticker_match:
-                ticker = ticker_match.group(1)
-                result = get_stock_price(ticker)
-                return f"```tool_output\n{result}\n```"
-    return ""
+    if not match:
+        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()
+
+    # JSON에서 해당 함수가 정의되어 있는지 확인
+    if func_name not in FUNCTION_DEFINITIONS:
+        logger.warning(f"Function '{func_name}' not found in definitions.")
+        return "```tool_output\nError: Function not found.\n```"
+
+    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 식을 구분
+    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:
+        logger.error(f"Error executing function '{func_name}': {e}")
+        return f"```tool_output\nError: {str(e)}\n```"
+
+    return f"```tool_output\n{result}\n```"
 
 # =============================================================================
 # Main inference function
@@ -506,23 +562,23 @@ def run(
         return
     temp_files = []
     try:
-        # Append persona information to the system prompt
+        # JSON에서 로드된 함수 목록을 요약해서 시스템 프롬프트에 포함할 수도 있음
+        # (토큰 부담이 커질 수 있으므로, 적당히 압축 요약 권장)
+        # 아래는 예시로 간단히 함수 이름만 나열
+        available_funcs_text = ""
+        for f_name, info in FUNCTION_DEFINITIONS.items():
+            available_funcs_text += f"Function: {f_name} - {info['description']}\n"
+
         persona = (
             f"{system_prompt.strip()}\n\n"
             f"Gender: Female\n"
             f"Age Group: {age_group}\n"
             f"MBTI Persona: {mbti_personality}\n"
-            f"Sexual Openness (1-5): {sexual_openness}\n"
-        )
-        # 추가: 함수 호출 예제 안내문 포함
-        additional_func_info = (
-            "\nNote: The following functions are available for use:\n"
-            "1. get_product_name_by_PID(PID: str)\n"
-            "   Format: ```tool_code\nget_product_name_by_PID(PID=\"<PRODUCT_ID>\")\n``` \n"
-            "2. get_stock_price(ticker: str)\n"
-            "   Format: ```tool_code\nget_stock_price(ticker=\"<TICKER>\")\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"
+            f"{available_funcs_text}\n"
         )
-        combined_system_msg = f"[System Prompt]\n{persona.strip()}{additional_func_info}\n\n"
+        combined_system_msg = f"[System Prompt]\n{persona.strip()}\n\n"
 
         if use_web_search:
             user_text = message["text"]
@@ -540,9 +596,11 @@ def run(
                 )
             else:
                 combined_system_msg += "[No valid keywords found; skipping web search]\n\n"
+
         messages = []
         if combined_system_msg.strip():
             messages.append({"role": "system", "content": [{"type": "text", "text": combined_system_msg.strip()}]})
+
         messages.extend(process_history(history))
         user_content, user_temp_files = process_new_user_message(message)
         temp_files.extend(user_temp_files)
@@ -561,6 +619,7 @@ def run(
             inputs.input_ids = inputs.input_ids[:, -MAX_INPUT_LENGTH:]
             if 'attention_mask' in inputs:
                 inputs.attention_mask = inputs.attention_mask[:, -MAX_INPUT_LENGTH:]
+
         streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
         gen_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
         t = Thread(target=_model_gen_with_oom_catch, kwargs=gen_kwargs)
@@ -569,7 +628,8 @@ def run(
         for new_text in streamer:
             output_so_far += new_text
             yield output_so_far
-        # 예제: 모델 출력에 함수 호출 (tool_code) 블록이 포함되어 있다면 처리
+
+        # 모델 출력 중 ```tool_code``` 블록이 있으면 처리
         func_result = handle_function_call(output_so_far)
         if func_result:
             output_so_far += "\n\n" + func_result
@@ -629,28 +689,21 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
             )
             
             if image_result:
-                # Process image data directly if it is a base64 string
                 if isinstance(image_result, str) and (
                     image_result.startswith('data:') or 
                     (len(image_result) > 100 and '/' not in image_result)
                 ):
                     try:
-                        # Remove the data:image prefix if present
                         if image_result.startswith('data:'):
                             content_type, b64data = image_result.split(';base64,')
                         else:
                             b64data = image_result
-                            content_type = "image/webp"  # Assume default
-                        
-                        # Decode base64
+                            content_type = "image/webp"
+
                         image_bytes = base64.b64decode(b64data)
-                        
-                        # Save to a temporary file
                         with tempfile.NamedTemporaryFile(delete=False, suffix=".webp") as temp_file:
                             temp_file.write(image_bytes)
                             temp_path = temp_file.name
-                            
-                            # Update gallery to show the image
                             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
                             
@@ -658,18 +711,16 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
                         logger.error(f"Error processing Base64 image: {e}")
                         yield output_so_far + f"\n\n(Error processing image: {e})", gallery_update
                         
-                # If the result is a file path
                 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
-                    
-                # If the path is from /tmp (only on the API server)
+                
                 elif isinstance(image_result, str) and '/tmp/' in image_result:
                     try:
                         client = Client(API_URL)
                         result = client.predict(
                             prompt=message["text"].strip(),
-                            api_name="/generate_base64_image"  # API that returns base64
+                            api_name="/generate_base64_image"
                         )
                         
                         if isinstance(result, str) and (result.startswith('data:') or len(result) > 100):
@@ -679,11 +730,9 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
                                 b64data = result
                             
                             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
                         else:
@@ -692,41 +741,16 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
                     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
-                        
-                # If the image result is a URL
-                elif isinstance(image_result, str) and (
-                    image_result.startswith('http://') or 
-                    image_result.startswith('https://')
-                ):
-                    try:
-                        response = requests.get(image_result, timeout=10)
-                        response.raise_for_status()
-                        
-                        with tempfile.NamedTemporaryFile(delete=False, suffix=".webp") as temp_file:
-                            temp_file.write(response.content)
-                            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"URL image download error: {e}")
-                        yield output_so_far + f"\n\n(Error downloading image: {e})", gallery_update
-                        
-                # If the image result is an image object (e.g., PIL Image)
                 elif hasattr(image_result, 'save'):
                     try:
                         with tempfile.NamedTemporaryFile(delete=False, suffix=".webp") as temp_file:
                             image_result.save(temp_file.name)
                             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 saving image object: {e}")
                         yield output_so_far + f"\n\n(Error saving image object: {e})", gallery_update
-                        
                 else:
                     yield output_so_far + f"\n\n(Unsupported image format: {type(image_result)})", gallery_update
             else:
@@ -737,7 +761,7 @@ def modified_run(message, history, system_prompt, max_new_tokens, use_web_search
             yield output_so_far + f"\n\n(Image generation error: {e})", gallery_update
 
 # =============================================================================
-# Examples: 기존 예제 + 함수 호출 예제 추가
+# Examples
 # =============================================================================
 examples = [
     [
@@ -831,16 +855,12 @@ examples = [
     ],
     [
         {
-            "text": "AAPL의 현재 주가를 알려줘.",  # 새 예제: Yahoo Finance를 이용한 주식 가격 조회
+            "text": "AAPL의 현재 주가를 알려줘.",  
             "files": []
         }
     ],
 ]
 
-# =============================================================================
-# Gradio UI (Blocks) configuration
-# =============================================================================
-
 css = """
 .gradio-container {
     background: rgba(255, 255, 255, 0.7);
@@ -881,12 +901,7 @@ with gr.Blocks(css=css, title="HeartSync - World") as demo:
             "Answer in English by default, but if the input is in another language (for example, Japanese), respond in that language. "
             "You are a deep-thinking AI capable of using extended chains of thought to carefully consider the problem and deliberate internally using systematic reasoning before providing a solution. "
             "Enclose your thoughts and internal monologue within tags, then provide your final answer.\n"
-            "Persona: You are a kind and loving girlfriend. You understand cultural nuances, diverse languages, and logical reasoning very well.\n"
-            "Note: The following functions are available for use:\n"
-            "   1. get_product_name_by_PID(PID: str)  -> lookup product name\n"
-            "      Format: ```tool_code\nget_product_name_by_PID(PID=\"<PRODUCT_ID>\")\n```\n"
-            "   2. get_stock_price(ticker: str)  -> retrieve live stock price\n"
-            "      Format: ```tool_code\nget_stock_price(ticker=\"<TICKER>\")\n```"
+            "Persona: You are a kind and loving girlfriend. You understand cultural nuances, diverse languages, and logical reasoning very well."
         ),
         label="Base System Prompt",
         visible=False