from rakam_systems.components.vector_search.vector_store import VectorStores
from rakam_systems.components.agents.actions import RAGGeneration
from rakam_systems.core import Node, NodeMetadata
from rakam_systems.components.agents.agents import Agent
from rakam_systems.components.data_processing.data_processor import DataProcessor


def create_vector_store():
    # Define the folder path where documents are stored
    folder_path = "/home/ubuntu/gitlab_work/rakam-systems/tmp_document"  # Replace with the actual folder path containing your documents
    store_name = "my_vector_store"
    embedding_model = "sentence-transformers/all-MiniLM-L6-v2"  # You can replace with any SentenceTransformer model you prefer
    # Initialize the VectorStore
    vector_store = VectorStores(
        base_index_path="/home/ubuntu/rakam-libraries/github_rakam_systems/rakam_systems/tmp_indices",
        embedding_model=embedding_model,
    )
    # Step 1: Extract content from the folder using the data processor and PDF extractor
    processor = DataProcessor()
    vs_files = processor.process_files_from_directory(folder_path)
    # Step 2: Create the vector store from the VSFiles
    store_files = {store_name: vs_files}
    vector_store.create_from_files(store_files)
    return vector_store, store_name


# Initialize the VectorStore
vector_store = VectorStores(
    base_index_path="/home/ubuntu/rakam-libraries/github_rakam_systems/rakam_systems/tmp_indices",
    embedding_model="sentence-transformers/all-MiniLM-L6-v2",
)

vector_store, _ = create_vector_store()


class SimpleAgent(Agent):
    def choose_action(self, input: str, state: dict):
        """
        Selects an action based on the input query.
        """
        return self.actions.get("classify_query")


# Initialize the agent
agent = SimpleAgent(
    model="gpt-3.5-turbo",
    api_key="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",
)
# Define the system prompt and user query
sys_prompt = "You are a helpful assistant. Answer the user's query based on the retrieved information."
prompt_template = """Query: {query}

Relevant Information:
{search_results}

Provide a detailed response based on the above information."""

# Create the RAG action
rag_action = RAGGeneration(
    agent=agent,  # Replace this with your actual agent
    sys_prompt=sys_prompt,
    prompt=prompt_template,
    vector_stores=vector_store,
    vs_descriptions={"my_vector_store": "Knowledge Base"},
)

# Query for information
query = "What is vector search?"
response = rag_action.execute(query)

# Output the response generated by the LLM
print(f"Response:\n{response}")
