LangChainとLlamaIndexでのWeb検索: ツールとretriever

import os

import requests
from langchain_core.tools import tool

API = "https://llmwebsearch.com/api/v1/search"
HEADERS = {"Authorization": "Bearer " + os.environ["LLMWEBSEARCH_API_KEY"]}


@tool
def search_web(query: str, read: int = 3) -> str:
    """Search the live web. Returns numbered sources with URLs; the first `read` include page text."""
    body = {"q": query, "max_results": 5, "read": max(0, min(read, 10)), "max_chars": 5000}
    resp = requests.post(API, headers=HEADERS, json=body, timeout=60)
    resp.raise_for_status()
    parts = []
    for i, r in enumerate(resp.json()["results"], start=1):
        text = r["content"] if r["content_status"] == "ok" else "(snippet only) " + r["snippet"]
        parts.append(f"[{i}] {r['title']} - {r['url']}\n{text}")
    return "\n\n".join(parts)


print(search_web.invoke({"query": "postgres connection pool size"})[:500])

import os
from typing import List

import requests
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever


class LLMWebSearchRetriever(BaseRetriever):
    """Live web search results as Documents, one per page that was read."""

    max_results: int = 5
    read: int = 3
    max_chars: int = 6000

    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
        resp = requests.post(
            "https://llmwebsearch.com/api/v1/search",
            headers={"Authorization": "Bearer " + os.environ["LLMWEBSEARCH_API_KEY"]},
            json={"q": query, "max_results": self.max_results, "read": self.read, "max_chars": self.max_chars},
            timeout=60,
        )
        resp.raise_for_status()
        docs = []
        for r in resp.json()["results"]:
            if r["content_status"] != "ok":
                continue
            meta = {"source": r["url"], "title": r["title"], "position": r["position"], "published_at": r.get("published_at")}
            docs.append(Document(page_content=r["content"], metadata=meta))
        return docs


retriever = LLMWebSearchRetriever(read=3)
for doc in retriever.invoke("pgbouncer transaction mode limitations"):
    print(doc.metadata["position"], doc.metadata["source"], len(doc.page_content))

import os

import requests
from llama_index.core.tools import FunctionTool


def search_web(query: str, read: int = 3) -> str:
    """Search the live web and return numbered sources; the first `read` include their page text."""
    resp = requests.post(
        "https://llmwebsearch.com/api/v1/search",
        headers={"Authorization": "Bearer " + os.environ["LLMWEBSEARCH_API_KEY"]},
        json={"q": query, "max_results": 5, "read": max(0, min(read, 10)), "max_chars": 5000},
        timeout=60,
    )
    resp.raise_for_status()
    parts = []
    for i, r in enumerate(resp.json()["results"], start=1):
        text = r["content"] if r["content_status"] == "ok" else "(snippet only) " + r["snippet"]
        parts.append(f"[{i}] {r['title']} - {r['url']}\n{text}")
    return "\n\n".join(parts)


search_tool = FunctionTool.from_defaults(fn=search_web)
print(search_tool.metadata.name)
print(search_tool.call(query="postgres connection pool size").content[:300])