search.py

from langchain_chroma import Chroma
from llm import embedding_model

vector_store = Chroma(
    embedding_function=embedding_model,
    collection_name="sample",
    persist_directory="./.chroma_db"
)

def search(query, top_k=3):
    """
    Perform a similarity search in the vector store.

    Args:
        query (str): The search query.
        top_k (int): The number of top similar documents to retrieve.
    """
    retriever = vector_store.as_retriever(search_kwargs={'k': top_k})
    docs = retriever.invoke(query)
    return docs

if __name__ == "__main__":
    query = "Explain Layout Analysis Model"
    results = search(query)
    print("Search Results:")
    for doc in results:
        print(f"Metadata: {doc.metadata}")
        print(f"- {doc.page_content}")