RAG/service/content_service.py

import json

import requests

from core.database import DBHelper
from data_models.contents import Contents


def get_contents(page_num, page_size, dataset_id):
    db_helper = DBHelper()
    res = db_helper.get_paginated(Contents, page_num, page_size, order_by={'id': 'desc'}, dataset_id=dataset_id,
                                  doc_status=1)
    data = []
    for entity in res["entities"]:
        data.append({'text': entity.text, 'title': entity.title, 'doc_id': entity.doc_id})
    res['entities'] = data
    return res


def get_content(doc_id):
    db_helper = DBHelper()
    content = db_helper.get(Contents, doc_id=doc_id)
    data = {'title': content.title, 'text': content.text, 'doc_id': content.doc_id}
    return data


def add_contents(content_param):
    try:
        response = requests.post(
            url='http://61.48.133.26:8001/api/chunk',
            json={
                "text": content_param.text,
                "title": content_param.title,
                "dataset_id": content_param.datasetId},
            headers={"Content-Type": "application/json"},
        )
        doc_id = response.json()['doc_id']
        if doc_id:
            return True
    except Exception as e:
        print(e)
    return False


if __name__ == '__main__':
    print(json.dumps(get_contents(1, 10), ensure_ascii=False))