# -*- coding: utf-8 -*- from pydantic import BaseModel, Field import requests class Tools: class Valves(BaseModel): URL: str = Field( default="localhost", description="The ragflow url", ) API_KEY: str = Field( default="apikey", description="The api key", ) DATASET_IDS: list = Field( default=["id"], description="The IDs of the datasets to search" ) PAGE: int = Field( default=1, description="Specifies the page on which the chunks will be displayed" ) PAGE_SIZE: int = Field( default=30, description="The maximum number of chunks on each page" ) SIMILARITY_THRESHOLD: float = Field( default=0.2, description="The minimum similarity score" ) VECTOR_SIMILARITY_WEIGHT: float = Field( default=0.3, description="The weight of vector cosine similarity" ) def __init__(self): self.valves = self.Valves() self.headers = { "Content-Type": "application/json", "Authorization": "Bearer {}".format(self.valves.API_KEY) } def retrieval(self, question: str) -> str: payload = { "question": question, "dataset_ids": self.valves.DATASET_IDS, "page": self.valves.PAGE, "page_size": self.valves.PAGE_SIZE, "similarity_threshold": self.valves.SIMILARITY_THRESHOLD, "vector_similarity_weight": self.valves.VECTOR_SIMILARITY_WEIGHT } contents = [] try: response = requests.post(self.valves.URL, headers=self.headers, json=payload) if response.status_code == 200: result = response.json() if "data" in result: data = result["data"] if "chunks" in data: chunks = data["chunks"] for chunk in chunks: contents.append(chunk["content"]) else: return "Data parse error" else: return "Data parse error" return "。".join(contents) else: return response.status_code except Exception as e: return "Error:{}".format(str(e))