neuronslabs/ComfyKnowledgeGraph

1from dotenv import load_dotenv2from llama_index.core import VectorStoreIndex3import os4from llama_index.llms.openai import OpenAI5from llama_index.core import StorageContext, Settings, load_index_from_storage6from llama_index.core import VectorStoreIndex, SimpleDirectoryReader7from retrieve import get_latest_dir8from datetime import datetime9 10 11load_dotenv()12 13Settings.llm = OpenAI(temperature=0, model="gpt-3.5-turbo")14Settings.chunk_size = 204815Settings.chunk_overlap = 2416 17 18def create_embedding():19    """20    Create an embedding from the given directory.21 22    Returns:23    VectorStoreIndex: The index of the embedding from docs in the directory.24    """25 26    output_dir = os.getenv("EMBEDDING_DIR")27    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")28    embedding_path = f"{output_dir}/{timestamp}"29 30    documents = SimpleDirectoryReader(os.getenv("PROD_SPEC_DIR")).load_data()31 32    index = VectorStoreIndex.from_documents(documents, show_progress=True)33    index.storage_context.persist(persist_dir=embedding_path)34 35    return index36 37 38def load_embedding():39    """40    Load the latest embedding from the directory.41 42    Returns:43    VectorStoreIndex: The index of the embedding from the latest directory.44    """45    PERSIST_DIR = get_latest_dir(os.getenv("EMBEDDING_DIR"))46    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)47    index = load_index_from_storage(storage_context)48    return index49 50 51def query_rag_qa(rag_index, query, search_level):52    """53    Query the RAG model for a given query.54 55    Args:56    rag_index (VectorStoreIndex): The RAG model index.57    query (str): The query to ask the RAG model.58    search_level (int): The max search level to use for the RAG model.59 60    Returns:61    tuple: The response, nodes, and reference text from the RAG model.62    """63    myretriever = rag_index.as_retriever(64        include_text=True,65        similarity_top_k=search_level,66    )67    query_engine = rag_index.as_query_engine(68        sub_retrievers=[69            myretriever,70        ],71        include_text=True,72        similarity_top_k=search_level,73    )74    response = query_engine.query(query)75    nodes = myretriever.retrieve(query)76 77    reference_text = []78    for node in nodes:79        reference_text.append(node.text)80 81    return response, nodes, reference_text82