main.py

import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from dotenv import load_dotenv
from pinecone_manager import PineconeManager
import uuid
import logging
from typing import Any, Dict, List, Optional
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
EMBEDDING_MODEL_NAME = os.getenv("EMBEDDING_MODEL_NAME")
PINECONE_INDEX = os.getenv("PINECONE_INDEX")

if not all([OPENAI_API_KEY, PINECONE_API_KEY, EMBEDDING_MODEL_NAME, PINECONE_INDEX]):
    raise ValueError("Missing keys. Please check your .env file.")

manager = PineconeManager(
    api_key=os.getenv("PINECONE_API_KEY"),
    index_name=os.getenv("PINECONE_INDEX")
)

embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL_NAME, openai_api_key=OPENAI_API_KEY)

def sanitize_namespace(filename):
    base = os.path.splitext(filename)[0]
    return base.replace(" ", "_").lower()


def process_file(file_path):
    loader = PyPDFLoader(file_path)
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000, chunk_overlap=200, add_start_index=True
    )
    all_splits = text_splitter.split_documents(docs)
    # embeddings = OpenAIEmbeddings(model=os.getenv("EMBEDDING_MODEL_NAME"), openai_api_key=os.getenv("OPENAI_API_KEY"))

    namespace = sanitize_namespace(os.path.basename(file_path))

    ''' when using PineconeVectorStore locally'''
    # vectorstore = PineconeVectorStore(index_name=PINECONE_INDEX, embedding=embeddings)
    # vectorstore.from_documents(all_splits, embeddings)

    ''' when using FAISS vectorstore locally'''
    # vectorstore = FAISS.from_documents(all_splits, embeddings)
    # vectorstore.save_local(index_name)
    # return pdf_name, embeddings

    manager.upsert_documents(
        documents=all_splits,
        embeddings=embeddings,
        namespace=namespace
    )
    return namespace


def query_index(query, namespace):
    ''' when using FAISS '''
    # vectorstore = FAISS.load_local(index_name, embeddings, allow_dangerous_deserialization=True)
    # results = vectorstore.similarity_search_with_score(query)
    # for doc, score in results:
    #     print(f"Score: {score}\n")
    #     print(doc)

    query_vector = embeddings.embed_query(query)
    results = manager.search(
        vector=query_vector,
        namespace=namespace,
        top_k=3
    )
    return {
        "results": [
            {
                "text": match.metadata["text"],
                "score": match.score
            } for match in results.matches
        ]
    }


def main():
    file_path = input("Enter the path to the PDF file: ")
    namespace = process_file(file_path)

    while True:
        namespace_search = input("Enter namespace (or 'exit' to quit): ")
        if namespace_search.lower() == 'exit':
            break
        query = input("Enter query (or 'exit' to quit): ")
        if query.lower() == 'exit':
            break
        res = query_index(query, namespace_search)
        print(res)





if __name__ == "__main__":
    main()