redis_ops.py

import json
from typing import Dict, List

import redis


class RedisOperations:
    def __init__(self, host='localhost', port=6379, db=0):
        self.client = redis.Redis(host=host, port=port, db=db)

    def store_document(self, content: str, chunks: List[str], embeddings: List[List[float]]):
        self.client.hset("document", mapping={
            "content": content,
            "chunks": json.dumps(chunks),
            "embeddings": json.dumps(embeddings)
        })

    def get_document(self) -> Dict[str, any]:
        doc = self.client.hgetall("document")
        if doc:
            return {
                "content": doc[b"content"].decode('utf-8'),
                "chunks": json.loads(doc[b"chunks"].decode('utf-8')),
                "embeddings": json.loads(doc[b"embeddings"].decode('utf-8'))
            }
        return {"content": "", "chunks": [], "embeddings": []}

    def clear_db(self):
        """Clear the stored document."""
        self.client.delete("document")
        print("Document cleared from Redis.")

    def print_db_contents(self) -> Dict[str, any]:
        """Retrieve and return the stored document."""
        doc = self.get_document()
        return {
            "document": doc["content"],
            "chunks_count": len(doc["chunks"]),
            "embeddings_count": len(doc["embeddings"])
        }