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
from data_models.dataset import Dataset


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


def get_embedding_content_data(query, dataset_ids):
    res = []
    db_helper = DBHelper()
    results = get_embedding_data(query, dataset_ids)
    if results:
        for result in results:
            content_chunk = db_helper.get(ContentChunks, doc_id=result['doc_id'], chunk_id=result['chunk_id'])
            dataset = db_helper.get(Dataset, id=content_chunk.dataset_id)
            dataset_name = None
            if dataset:
                dataset_name = dataset.name
            res.append(
                {'docId': content_chunk.doc_id, 'content': content_chunk.text,
                 'contentSummary': content_chunk.summary, 'score': result['score'], 'datasetName': dataset_name})
    return res