Upload 5 files

agent_for_unit4/__init__.py

@@ -0,0 +1,3 @@
+from .agent import manager_agent, prepare_for_input
+
+__all__ = ["manager_agent", "prepare_for_input"]

agent_for_unit4/agent.py

@@ -0,0 +1,113 @@
+import base64
+import os
+from pathlib import Path
+from textwrap import dedent
+
+from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMModel, VisitWebpageTool
+
+from .tools import RetrieveCSVStorageTool, SpeechRecognitionTool, VisualQATool, WikiTool, fetch_text_content, read_excel
+
+
+def configure_open_telemetry() -> None:
+    try:
+        from openinference.instrumentation.smolagents import SmolagentsInstrumentor
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+    except ImportError:
+        print("OpenTelemetry packages are not installed. Please install them to enable tracing.")
+        return None
+
+    try:
+        langfuse_public_key = os.environ["LANGFUSE_PUBLIC_KEY"]
+        langfuse_secret_key = os.environ["LANGFUSE_SECRET_KEY"]
+    except KeyError:
+        print("LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY must be set in the environment variables.")
+        return None
+
+    LANGFUSE_AUTH = base64.b64encode(f"{langfuse_public_key}:{langfuse_secret_key}".encode()).decode()
+    os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = "https://cloud.langfuse.com/api/public/otel"
+    os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"Authorization=Basic {LANGFUSE_AUTH}"
+
+    trace_provider = TracerProvider()
+    trace_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter()))
+
+    SmolagentsInstrumentor().instrument(tracer_provider=trace_provider)
+
+
+configure_open_telemetry()
+
+wiki_storage_tool = RetrieveCSVStorageTool(
+    table_name="wiki",
+    init_storage=True,
+    storage_path="./storage",
+)
+
+wiki_agent = CodeAgent(
+    name="wiki_agent",
+    description= """A wiki agent that can search and retrieve information from Wikipedia.
+    It is specialized for handling wikipedia articles, and is recommended over web_agent for retrieving information from wikipedia.""",
+    model=LiteLLMModel(model_id="openrouter/qwen/qwen-2.5-coder-32b-instruct"),
+    tools=[
+        DuckDuckGoSearchTool(),
+        wiki_storage_tool,
+        WikiTool(storage=wiki_storage_tool.get_storage()),
+    ],
+    max_steps=10,
+    additional_authorized_imports=["pandas"],
+)
+
+
+web_agent = CodeAgent(
+    name="web_agent",
+    description="A web agent that can search and visit webpages.",
+    model=LiteLLMModel(model_id="openrouter/qwen/qwen-2.5-coder-32b-instruct"),
+    tools=[
+        DuckDuckGoSearchTool(max_results=10),
+        VisitWebpageTool(),
+    ],
+    verbosity_level=2,
+    max_steps=10,
+)
+
+
+manager_agent = CodeAgent(
+    name = "manager_agent",
+    model=LiteLLMModel(
+        model_id="openrouter/qwen/qwq-32b",
+    ),
+    tools=[
+        fetch_text_content,  # fetch text content from a URL
+        SpeechRecognitionTool(),  # Audio to text
+        VisualQATool(),  # Visual Question Answering
+        read_excel,  # Read Excel files
+    ],
+    managed_agents=[
+        wiki_agent,
+        web_agent,
+    ],
+    additional_authorized_imports=["pandas", "requests"],
+    planning_interval=5,
+    verbosity_level=2,
+    max_steps=15,
+)
+
+
+def parse_file_name(file_base_url: str, file_name: str) -> str:
+    if file_name == "":
+        return "not provided"
+    return file_base_url + Path(file_name).stem
+
+
+def prepare_for_input(question: dict, file_base_url: str) -> str:
+    input_text = dedent(f"""\
+    Question:
+    {question["question"]}
+
+    If necessary, use the following file (they may not be provided)
+    file_type: {Path(question["file_name"]).suffix}
+    file: {parse_file_name(file_base_url, question["file_name"])}
+
+    Video analysis tools are currently unavailable.
+    If the question is about analyzing the video (e.g. questions about Youtube link and mp4), answer 'No Answer'.""")
+    return input_text

agent_for_unit4/db.py

@@ -0,0 +1,49 @@
+import shelve
+import shutil
+from pathlib import Path
+from typing import Any, Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class ShelveDB(Generic[T]):
+    dir_path: Path
+
+    def __init__(self, db_name: str, init: bool) -> None:
+        self.db_path = self.dir_path / db_name
+
+        if init:
+            self.dir_path.mkdir(parents=True, exist_ok=True)
+            for file_path in self.dir_path.glob(f"{db_name}*"):
+                if file_path.is_file():
+                    file_path.unlink()
+                elif file_path.is_dir():
+                    shutil.rmtree(file_path)
+
+    @classmethod
+    def from_table(cls, table: str) -> "ShelveDB":
+        return cls(table, False)
+
+    def save(self, key: str, value: Any) -> None:
+        with shelve.open(str(self.db_path)) as db:
+            db[key] = value
+
+    def fetch(self, key: str) -> T | None:
+        with shelve.open(str(self.db_path)) as db:
+            return db.get(key, None)
+
+    def delete(self, key: str) -> bool:
+        with shelve.open(str(self.db_path)) as db:
+            if key in db:
+                del db[key]
+                return True
+            return False
+        
+    def clear(self) -> None:
+        with shelve.open(str(self.db_path)) as db:
+            for key in list(db.keys()):
+                del db[key]
+
+    def list_keys(self) -> list[str]:
+        with shelve.open(str(self.db_path)) as db:
+            return list(db.keys())

agent_for_unit4/tools.py

@@ -0,0 +1,205 @@
+import os
+from io import BytesIO
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+import requests
+from huggingface_hub import InferenceClient
+from smolagents import Tool, tool
+
+from .db import ShelveDB
+from .wiki import get_wiki_content
+
+
+### convert table to markdown
+@tool
+def convert_pandas_table_to_markdown(table: pd.DataFrame) -> str:
+    """
+    Converts a pandas DataFrame to a markdown table.
+
+    Args:
+        table (pd.DataFrame): The DataFrame to convert.
+
+    Returns:
+        str: The markdown representation of the table.
+    """
+    return str(table.to_markdown())
+
+
+### fetch text tool
+@tool
+def fetch_text_content(url: str) -> str:
+    """
+    Fetches the text content from a given URL.
+
+    Args:
+        url (str): The URL to fetch the text from.
+
+    Returns:
+        str: The text content of the page.
+    """
+    try:
+        response = requests.get(url, timeout=30)
+        response.raise_for_status()  # Raise an error for bad responses
+        return response.text
+    except requests.RequestException as e:
+        return f"Error fetching URL: {e}"
+
+
+### Storage Tool
+class RetrieveCSVStorageTool(Tool):
+    name = "retrieve_csv_storage_tool"
+    description = "Retrieves a CSV file from the storage and returns it as a pandas DataFrame."
+    inputs = {
+        "key": {
+            "type": "string",
+            "description": "The key to retrieve data from the table.",
+        },
+    }
+    output_type = "any"
+
+    def __init__(self, table_name: str, init_storage: bool, storage_path: str | None = None, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        if storage_path is not None:
+            ShelveDB.dir_path = Path(storage_path)
+        self.storage = ShelveDB[pd.DataFrame](table_name, init=init_storage)
+
+    def get_storage(self) -> ShelveDB[pd.DataFrame]:
+        return self.storage
+
+    def forward(self, key: str) -> pd.DataFrame:
+        try:
+            # Retrieve the CSV file from storage
+            dataframe = self.storage.fetch(key)
+        except Exception as e:
+            return f"Error retrieving data: {e}"
+        else:
+            if dataframe is None:
+                raise ValueError(f"No data found for key: {key}")
+            return dataframe
+
+
+### Wikipedia Content Extraction Tool
+
+
+class WikiTool(Tool):
+    name = "wiki_tool"
+    description = """Get Wikipedia page content and tables.
+    Returns a tuple containing the page content and a dictionary of tables extracted from the page.
+    The page content is prefixed with the retrieved table key ({{table_1}}, {{table_2}}, ...).
+    To understand what is contained in the tables, it is recommended to first display the content.
+    Example 1:
+        content, tables = get_wiki_content("Python_(programming_language)")
+        print(content)
+    
+    The retrieved table object is are stored in storage.
+    They can be retrieved using "retrieve_csv_storage_tool".
+    Example 2:
+        table:pd.DataFrame = retrieve_csv_storage_tool("table_1")
+    """
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The title of the Wikipedia page to visit. For example, 'Python_(programming_language)'.",
+        },
+        "language": {
+            "type": "string",
+            "description": "The language of the Wikipedia page. For example, 'en' for English, 'ja' for Japanese.",
+        },
+    }
+    output_type = "array"
+
+    def __init__(self, storage: ShelveDB[Any], *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.storage = storage
+
+    def forward(self, query: str, language: str) -> tuple[str, dict[str, pd.DataFrame]]:
+        content, tables = get_wiki_content(query, language)
+        self.storage.clear()
+        for table_key, df in tables.items():
+            self.storage.save(table_key, df)
+        return content, tables
+
+
+### Visual Question Answering Tool
+
+
+def request_visual_qa(client: InferenceClient, question: str, image_url: str) -> str:
+    contents = [{"type": "text", "text": question}, {"type": "image_url", "image_url": {"url": image_url}}]
+    res = client.chat_completion(messages=[{"role": "user", "content": contents}], model="qwen/qwen2.5-vl-32b-instruct")
+    content = res.choices[0].message.content
+    if content is None:
+        raise ValueError("No content returned from the model.")
+    return content
+
+
+class VisualQATool(Tool):
+    name = "visual_qa_tool"
+    description = "A tool that can answer questions about image."
+    inputs = {
+        "image_url": {
+            "type": "string",
+            "description": "The URL of the image to analyze. No extension needed.",
+        },
+        "question": {
+            "type": "string",
+            "description": "The question to ask about the image.",
+        },
+    }
+    output_type = "string"
+    client = InferenceClient(
+        base_url="https://openrouter.ai/api/v1",
+        api_key=os.environ["OPENROUTER_API_KEY"],
+    )
+
+    def forward(self, image_url: str, question: str) -> str:
+        try:
+            answer = request_visual_qa(self.client, question, image_url)
+        except Exception as e:
+            return f"Error: {str(e)}"
+        else:
+            return answer
+
+
+### Speech Recognition Tool
+
+
+def request_speech_recognition(client: InferenceClient, audio_file: str, model: str = "openai/whisper-large-v3") -> str:
+    output = client.automatic_speech_recognition(audio_file, model=model)
+    return output.text
+
+
+class SpeechRecognitionTool(Tool):
+    name = "speech_recognition"
+    description = "Converts audio contents to text"
+    inputs = {"audio_url": {"type": "string", "description": "URL of the audio file to transcribe. No extension needed."}}
+    output_type = "string"
+    client = InferenceClient(provider="fal-ai")
+    _model = "openai/whisper-large-v3"
+
+    def forward(self, audio_url: str) -> str:
+        try:
+            transcription = request_speech_recognition(self.client, audio_url, model=self._model)
+        except Exception as e:
+            return f"Error: {str(e)}"
+        else:
+            return transcription
+
+
+### Excel Tool
+@tool
+def read_excel(file_url: str) -> pd.DataFrame:
+    """
+    Reads an Excel file from a given URL and returns the data as a DataFrame.
+
+    Args:
+        file_url (str): URL of the Excel file to read. No extension needed.
+    Returns:
+        pd.DataFrame: DataFrame containing the data from the first sheet of the Excel file
+    """
+    res = requests.get(file_url, timeout=30)
+    res.raise_for_status()
+    excel_data = BytesIO(res.content)
+    df = pd.read_excel(excel_data)
+    return df

agent_for_unit4/wiki.py

@@ -0,0 +1,180 @@
+import re
+from io import StringIO
+from typing import Any
+
+import pandas as pd
+import requests
+from bs4 import BeautifulSoup
+
+
+def process_list_element(list_element: Any, indent: int = 0) -> str:
+    """リスト要素を再帰的に処理する関数"""
+    result = []
+
+    is_ordered = list_element.name == "ol"
+
+    for i, li in enumerate(list_element.find_all("li", recursive=False)):
+        # リスト項目のテキストを取得
+        # ネストされたリストを除いたテキストを取得
+        item_text = ""
+        for content in li.contents:
+            if content.name not in ["ul", "ol"]:
+                item_text += str(content)
+
+        item_text = BeautifulSoup(item_text, "html.parser").get_text().strip()
+
+        # 順序付きリストなら番号を、そうでなければ記号を使用
+        prefix = "  " * indent + (f"{i + 1}. " if is_ordered else "* ")
+        if item_text:
+            result.append(prefix + item_text)
+
+        # ネストされたリストを処理
+        for nested_list in li.find_all(["ul", "ol"], recursive=False):
+            nested_content = process_list_element(nested_list, indent + 1)
+            if nested_content:
+                result.append(nested_content)
+
+    return "\n".join(result)
+
+
+def get_wiki_content(title: str, language: str = "en") -> tuple[str, dict[str, pd.DataFrame]]:
+    """
+    Get Wikipedia page content and tables.
+
+    Returns:
+        A tuple containing the page content as a string and a dictionary of tables
+        extracted from the page. The keys of the dictionary are "table_1", "table_2", etc.
+        and the values are pandas DataFrames representing the tables.
+
+    Example:
+        content, tables = get_wiki_content("Python_(programming_language)")
+        print(content)
+        print(tables["table_1"])  # Access the first table
+
+    Args:
+        title: wikipedia page title (e.g., "Python_(programming_language)")
+        language: wikipedia language (e.g., "en" for English, "ja" for Japanese)
+    """
+    # パースAPIのURLを構築
+    api_url = f"https://{language}.wikipedia.org/w/api.php"
+
+    # APIパラメータ
+    params = {
+        "action": "parse",
+        "page": title,
+        "format": "json",
+        "prop": "text",
+        "disabletoc": True,
+    }
+
+    # リクエストを送信
+    response = requests.get(api_url, params=params, timeout=30)  # type: ignore
+
+    # レスポンスをチェック
+    if response.status_code != 200:
+        raise Exception(f"api error: {response.status_code} - {response.text}")
+
+    # JSONレスポンスをパース
+    data = response.json()
+
+    # エラーチェック
+    if "error" in data:
+        raise Exception(f"api error: {data['error']['info']}")
+
+    if "parse" not in data:
+        raise Exception("api error: No parse data found")
+
+    # HTMLコンテンツを取得
+    html_content = data["parse"]["text"]["*"]
+
+    # HTMLをパース
+    soup = BeautifulSoup(html_content, "html.parser")
+    content_soup = BeautifulSoup(html_content, "html.parser")
+
+    # テーブル情報を取得
+    tables_dict: dict[str, pd.DataFrame] = {}
+    table_ids: list[tuple[str, str]] = []  # (table_id, table_html) のリスト
+
+    # ターゲットとするテーブルを特定: wikitableとinfobox
+    table_index = 1
+
+    # まず、infobox（バイオグラフィーテーブル）を処理
+    infoboxes = soup.find_all("table", class_=lambda c: c and "infobox" in c)
+    for i, table in enumerate(infoboxes):
+        table_id = f"table_{table_index}"
+        table_ids.append((table_id, str(table)))
+        table_index += 1
+
+    # 次に、wikitableを処理
+    wikitables = soup.find_all("table", class_="wikitable")
+    for i, table in enumerate(wikitables):
+        table_id = f"table_{table_index}"
+        table_ids.append((table_id, str(table)))
+        table_index += 1
+
+    # 抽出したテーブルをpandasで処理
+    for table_id, table_html in table_ids:
+        try:
+            dfs = pd.read_html(StringIO(table_html))
+            if dfs:
+                tables_dict[table_id] = dfs[0]
+        except Exception:
+            # テーブル解析に失敗した場合はスキップ
+            continue
+
+    # コンテンツ内のテーブルをプレースホルダに置き換え
+    table_placeholders: dict[str, str] = {}
+
+    # infoboxの処理
+    for i, table in enumerate(content_soup.find_all("table", class_=lambda c: c and "infobox" in c)):
+        table_id = f"table_{i + 1}"
+        if table_id in tables_dict:
+            placeholder = f"{{{{{table_id}}}}}"
+            table_placeholders[table_id] = placeholder
+            table_placeholder_tag = content_soup.new_tag("p")
+            table_placeholder_tag.string = placeholder
+            table.replace_with(table_placeholder_tag)
+
+    # wikitableの処理（インデックスは続きから）
+    wikitable_start_index = len(infoboxes) + 1
+    for i, table in enumerate(content_soup.find_all("table", class_="wikitable")):
+        table_id = f"table_{wikitable_start_index + i}"
+        if table_id in tables_dict:
+            placeholder = f"{{{{{table_id}}}}}"
+            table_placeholders[table_id] = placeholder
+            table_placeholder_tag = content_soup.new_tag("p")
+            table_placeholder_tag.string = placeholder
+            table.replace_with(table_placeholder_tag)
+
+    # クリーンな本文テキストを抽出
+    for element in content_soup.find_all(["sup", "div.hatnote", "div.navbox", "span.mw-editsection"]):
+        element.decompose()
+
+    # 見出し、パラグラフ、リストを取得
+    elements = content_soup.find_all(["h1", "h2", "h3", "h4", "h5", "h6", "p", "ul", "ol"])
+    text_content = []
+
+    for element in elements:
+        if element.name and element.name.startswith("h"):  # type: ignore
+            level = int(element.name[1])  # type: ignore
+            heading_text = element.get_text().strip()
+            if heading_text:  # 空の見出しをスキップ
+                text_content.append("\n" + "#" * level + " " + heading_text)
+        elif element.name == "p":  # type: ignore
+            paragraph_text = element.get_text().strip()
+            if paragraph_text:  # 空のパラグラフをスキップ
+                # テーブルプレースホルダの場合はそのまま追加
+                if re.match(r"^\{\{table_\d+\}\}$", paragraph_text):
+                    text_content.append(paragraph_text)
+                else:
+                    text_content.append(paragraph_text)
+        elif element.name in ["ul", "ol"] and element.parent.name not in ["li", "ul", "ol"]:  # type: ignore
+            # トップレベルのリストのみ処理（ネストされたものは親liで処理）
+            list_content = process_list_element(element)
+            if list_content:
+                text_content.append(list_content)
+
+    # テキストコンテンツを結合
+    content = "\n\n".join(text_content)
+
+    return content, tables_dict