openai_ef = embedding_functions.OpenAIEmbeddingFunction(
api_key=os.environ["OPENAI_API_KEY"], model_name="text-embedding-3-small"
)
chroma_client = chromadb.Client()
chroma_collection = chroma_client.get_or_create_collection(
"agentic-rag-demo-company-policies", embedding_function=openai_ef
)
chroma_collection.add(
ids=["1", "2", "3"],
documents=[
"The travel policy is: Employees must book travel through the company portal. Economy class flights and standard hotel rooms are covered. Meals during travel are reimbursed up to $75/day. All expenses require receipts.",
"The pto policy is: Full-time employees receive 20 days of paid time off per year, accrued monthly. PTO requests must be submitted at least 2 weeks in advance through the HR portal. Unused PTO can carry over up to 5 days into the next year.",
"The dress code is: Business casual attire is required in the office. This includes collared shirts, slacks or knee-length skirts, and closed-toe shoes. Jeans are permitted on Fridays. No athletic wear or overly casual clothing.",
],
)
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
chroma_index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
chroma_engine_policy = chroma_index.as_query_engine(similarity_top_k=1)