Update tavily_search.py

tavily_search.py

@@ -1,18 +1,17 @@
 import requests
 import gradio as gr
 import os
-import openai
-import time
 from urllib.parse import urlparse
 from requests.adapters import HTTPAdapter
 from requests.packages.urllib3.util.retry import Retry
+from groq import Groq  # Groqクライアントをインポート
+
+# Groq APIの設定
+client = Groq(api_key=os.getenv("GROQ_API_KEY"))
 
 # Tavily APIの設定
 TAVILY_API_URL = "https://api.tavily.com/search"
 API_KEY = os.getenv('TAVILY_API_KEY')
-OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
-
-openai.api_key = OPENAI_API_KEY
 
 def filter_results_by_domain(results, allowed_domains):
     filtered_results = []
@@ -24,12 +23,8 @@ def filter_results_by_domain(results, allowed_domains):
 
 def search_tavily(query, domain_filter, retries=3, backoff_factor=1.0, timeout=30):
     headers = {"Content-Type": "application/json"}
-
-    # キーワードが5文字以下の場合、デフォルトのフレーズを追加
     if len(query) <= 5:
         query += " example"
-
-    # 検索パラメータの定義
     params = {
         'api_key': API_KEY,
         'query': query,
@@ -37,17 +32,10 @@ def search_tavily(query, domain_filter, retries=3, backoff_factor=1.0, timeout=3
         'detail_level': 'high',
         'search_depth': 'advanced'
     }
-
-    retry_strategy = Retry(
-        total=retries,
-        status_forcelist=[429, 500, 502, 503, 504],
-        allowed_methods=["POST"],
-        backoff_factor=backoff_factor
-    )
+    retry_strategy = Retry(total=retries, status_forcelist=[429, 500, 502, 503, 504], allowed_methods=["POST"], backoff_factor=backoff_factor)
     adapter = HTTPAdapter(max_retries=retry_strategy)
     http = requests.Session()
     http.mount("https://", adapter)
-
     for attempt in range(retries):
         try:
             response = http.post(TAVILY_API_URL, headers=headers, json=params, timeout=timeout)
@@ -60,60 +48,30 @@ def search_tavily(query, domain_filter, retries=3, backoff_factor=1.0, timeout=3
                 return {'success': True, 'results': results['results']}
             elif response.status_code >= 500:
                 time.sleep(backoff_factor * (2 ** attempt))
-            else:
-                return {"success": False, "error": f"結果取得に失敗しました: {response.status_code}", 'results': []}
         except requests.exceptions.RequestException as e:
             time.sleep(backoff_factor * (2 ** attempt))
-
     return {"success": False, "error": "結果取得に失敗しました。後でもう一度試してください。", 'results': []}
 
 def summarize_results(results):
     if not results.get('success', False):
         return results.get('error', 'Unknown error occurred.')
-
     result_items = results.get('results', [])
     if not result_items:
         return "Error: No results found or unexpected data format."
-
-    snippets = []
-    references = []
-    for index, result in enumerate(result_items, start=1):
-        if 'content' in result and 'url' in result:
-            snippets.append(f"{result['content']} [^{index}]")
-            references.append(f"[^{index}]: {result['url']}")
-
+    snippets = [f"{result['content']} [^{index}]" for index, result in enumerate(result_items, start=1)]
+    references = [f"[^{index}]: {result['url']}" for index, result in enumerate(result_items, start=1)]
     joined_snippets = "\n\n".join(snippets)
     joined_references = "\n".join(references)
-
-    system_message = {
-        "role": "system",
-        "content": "あなたはプロのウェブリサーチャーです。次の検索結果を日本語でまとめ、詳細に説明してください。"
-    }
-    user_message = {
-        "role": "user",
-        "content": (
-            f"次のテキストに関する詳細な説明を、直接的なコピーまたは近いフレーズを避けてオリジナルの文章で作成し、『箇所書き』および『表』を適宜使用して日本語で説明してください。"
-            f"具体例を挙げる際は必ず表やグラフを用いてわかりやすく説明してください。"
-            f"テキストを生成する際に、参照した情報のURLを脚注やリンクとしてMarkdown形式で含めてください。\n\n{joined_snippets}\n\n{joined_references}"
-        )
-    }
-
-    try:
-        response = openai.ChatCompletion.create(
-            model="gpt-4o",
-            messages=[system_message, user_message],
-            temperature=0.7,
-            max_tokens=3000,
-            timeout=120
-        )
-        summary = response.choices[0]["message"]["content"].strip()
-        markdown_result = f"\n{summary}\n\n参照URL:\n{joined_references}"
-        html_result = markdown_result.replace("\n", "<br>").replace(" ", "&nbsp;")  # 簡単なHTML変換
-        return {'markdown': markdown_result, 'html': html_result}
-    except openai.error.OpenAIError as e:
-        return {"markdown": "Error: Unable to summarize the results.", "html": "Error: Unable to summarize the results."}
-    except requests.exceptions.Timeout as e:
-        return {"markdown": "Error: OpenAI API request timed out.", "html": "Error: OpenAI API request timed out."}
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {"role": "user", "content": "Translate the following text to Japanese:\n\n" + joined_snippets + "\n\n" + joined_references}
+        ],
+        model="llama3-8b-8192",
+        temperature=0.5  # 仮にtemperatureを設定できる場合
+    )
+    markdown_result = f"\n{chat_completion.choices[0].message.content}\n\n参照URL:\n{joined_references}"
+    html_result = markdown_result.replace("\n", "<br>").replace(" ", "&nbsp;")
+    return {'markdown': markdown_result, 'html': html_result}
 
 # Gradioインターフェースの設定
 def tavily_search_interface(selected_text, domain_filter):
@@ -129,7 +87,7 @@ def create_tavily_search_ui():
             search_button = gr.Button("検索", elem_classes="right-align")
             search_results_md = gr.Markdown(elem_classes="right-align")
             search_results_html = gr.HTML(elem_classes="right-align", visible=False)
-
             search_button.click(tavily_search_interface, inputs=[selected_text, domain_filter], outputs=[search_results_md, search_results_html])
+    return demo
 
-    return demo
+create_tavily_search_ui().launch()