RAG/utils/embedding_utils.py

import json

import requests
from core.config import logger
from core.database import DBHelper
from data_models.content_chunks import ContentChunks


def get_embedding_data(query):
    try:
        response = requests.post(
            url='http://61.48.133.26:8001/api/search',
            json={
                "query": query,
                "search_type": "by_vector",
                "limit": 5},
            headers={"Content-Type": "application/json"},
        )
        return response.json()['results']
    except Exception as e:
        logger.error(e)
    return []


def get_embedding_content_data(query):
    res = []
    db_helper = DBHelper()
    results = get_embedding_data(query)
    if results:
        for result in results:
            content_chunk = db_helper.get(ContentChunks, doc_id=result['doc_id'], chunk_id=result['chunk_id'])
            res.append(
                {'content': content_chunk.text, 'content_summary': content_chunk.summary, 'score': result['score']})
    return res


if __name__ == '__main__':
    results = get_embedding_content_data("帮我查询一些篮球相关的知识")
    print(results)