working cmd

document_handler.py

@@ -0,0 +1,22 @@
+import os
+import pickle
+from langchain_community.document_loaders.sitemap import SitemapLoader
+
+docs_file_path = 'sitemap_docs.pkl'
+
+def save_documents_to_disk(docs, file_path):
+    """Save the documents to a file using pickle."""
+    with open(file_path, 'wb') as file:
+        pickle.dump(docs, file)
+
+def load_documents_from_disk(file_path):
+    """Load the documents from a file if it exists."""
+    if os.path.exists(file_path):
+        with open(file_path, 'rb') as file:
+            return pickle.load(file)
+    return None
+
+def load_documents_from_sitemap(sitemap_url):
+    """Load documents from a sitemap URL using SitemapLoader."""
+    sitemap_loader = SitemapLoader(web_path=sitemap_url)
+    return sitemap_loader.load()

main.py

@@ -0,0 +1,38 @@
+import argparse
+from document_handler import load_documents_from_disk, load_documents_from_sitemap, save_documents_to_disk
+from vectorstore_handler import load_or_create_vectorstore, get_embeddings
+from query_executor import execute_query
+
+# Argument parsing for command line
+def main():
+    parser = argparse.ArgumentParser(description="Process a query to the Chroma vectorstore.")
+    parser.add_argument('query', type=str, help='The query to search in the vector store')
+    args = parser.parse_args()
+
+    # Path to save the documents
+    sitemap_url = "https://www.originws.it/page-sitemap.xml"
+    docs_file_path = 'sitemap_docs.pkl'
+
+    # Try to load documents from disk
+    docs = load_documents_from_disk(docs_file_path)
+
+    if docs is None:
+        print("Documents not found on disk, loading from sitemap...")
+        # Load documents using SitemapLoader
+        docs = load_documents_from_sitemap(sitemap_url)
+        save_documents_to_disk(docs, docs_file_path)
+        print("Documents saved to disk.")
+    else:
+        print("Documents loaded from disk.")
+
+    # Get embeddings and load/create the vectorstore
+    embeddings = get_embeddings()
+    vectorstore = load_or_create_vectorstore(docs, embeddings)
+
+    # Now that the vectorstore is ready, let's query it
+    question = args.query
+    response = execute_query(question, vectorstore)
+    print(response)
+
+if __name__ == "__main__":
+    main()

query_executor.py

@@ -0,0 +1,43 @@
+from langchain_openai import ChatOpenAI
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnablePassthrough
+
+RAG_TEMPLATE = """
+You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
+
+<context>
+{context}
+</context>
+
+Answer the following question:
+
+{question}"""
+
+def format_docs(docs):
+    """Format documents into a single string."""
+    return "\n\n".join(doc.page_content for doc in docs)
+
+def execute_query(question, vectorstore):
+    """Run the query against the vectorstore and return a response."""
+    print(f"Searching for: {question}")
+    docs = vectorstore.similarity_search(question, k=10)
+    print(f"Found {len(docs)} relevant documents for the query.")
+
+    # Set up the LLM and prompt handling
+    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+
+    # Define the RAG prompt template
+    rag_prompt = ChatPromptTemplate.from_template(RAG_TEMPLATE)
+
+    # Create the chain
+    chain = (
+        RunnablePassthrough.assign(context=lambda input: format_docs(input["context"]))
+        | rag_prompt
+        | llm
+        | StrOutputParser()
+    )
+
+    # Run the chain with the query
+    response = chain.invoke({"context": docs, "question": question})
+    return response

vectorstore_handler.py

@@ -0,0 +1,27 @@
+import os
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_openai import OpenAIEmbeddings
+from langchain_chroma import Chroma
+
+chroma_db_dir = 'chroma_vectorstore'
+
+def get_embeddings():
+    """Initialize and return OpenAI embeddings."""
+    return OpenAIEmbeddings(model="text-embedding-3-large")
+
+def load_or_create_vectorstore(docs, embeddings):
+    """Load or create a Chroma vectorstore."""
+    if os.path.exists(chroma_db_dir):
+        print("Loading existing Chroma vector store from disk...")
+        return Chroma(persist_directory=chroma_db_dir, embedding_function=embeddings)
+    
+    # Split documents if vectorstore doesn't exist
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)
+    all_splits = text_splitter.split_documents(docs)
+    print(f"Documents are split into {len(all_splits)} chunks from {len(docs)} documents.")
+    
+    # Create new vectorstore
+    print("Creating new Chroma vector store...")
+    vectorstore = Chroma.from_documents(documents=all_splits, embedding=embeddings, persist_directory=chroma_db_dir)
+    print(f"Vectorstore created and saved to {chroma_db_dir}")
+    return vectorstore