API: retrieval api

[Valdanitooooo]

What problem does this PR solve?

Add retrieval api on a specific knowledge base

#1102

Type of change

• New Feature (non-breaking change which adds functionality)

[Valdanitooooo]

Now, you can use RagFlow's retrieval API as a tool in your LLM Agent. 😉

import json
import os

import requests
from dotenv import load_dotenv

load_dotenv()

RAGFLOW_API_KEY = os.environ.get("RAGFLOW_API_KEY", "")
RAGFLOW_BASE_URL = os.environ.get("RAGFLOW_BASE_URL", "")

headers = {
    'Authorization': f'Bearer {RAGFLOW_API_KEY}',
    'Content-Type': 'application/json',
}

def ragflow_api_retrieval(kb_id, doc_count, prompt):
    payload = {
        "similarity_threshold": 0.2,
        "vector_similarity_weight": 0.3,
        "question": prompt,
        "doc_ids": [],
        "kb_id": kb_id,
        "page": 1,
        "size": doc_count
    }
    response = requests.post(
        RAGFLOW_BASE_URL + "/api/retrieval",
        headers={**headers},
        data=json.dumps(payload)
    )
    if response.status_code == 200:
        try:
            data = response.json()
            # print(data)
            chunks = data['data']['chunks']
            if len(chunks) > 0:
                docs_str = ""
                for chunk in chunks:
                    docs_str += chunk["content_ltks"].replace("\r", "\n") + "\n"
            return docs_str
        except Exception as e:
            print(str(e))
    else:
        print(f"Request failed with status code {response.status_code}")
    return "No data found"