[WIP]

.gitignore

@@ -1 +1,2 @@
-local setting.md
+local setting.md
+venv

app.py

@@ -1,5 +1,36 @@
 import streamlit as st
+from .chroma import searchengine
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+def main():
+    st.title("ChromaDB Search Engine")
+
+    model_name = "your_model_name"  # Replace with the actual model name
+    collection_name = "your_collection_name"  # Replace with the actual collection name
+
+    search_engine = searchengine(model_name, collection_name)
+
+    st.sidebar.header("Add Document")
+    text_input = st.sidebar.text_area("Enter Document Text")
+    metadata_input = st.sidebar.text_input("Enter Metadata")
+    add_button = st.sidebar.button("Add Document")
+
+    if add_button:
+        document_id = search_engine.count() + 1  # Incremental ID
+        search_engine.add(text_input, metadata_input, document_id)
+        st.sidebar.success(f"Document added with ID: {document_id}")
+
+    st.sidebar.header("Search")
+    query = st.sidebar.text_input("Enter Search Query")
+    search_button = st.sidebar.button("Search")
+
+    if search_button:
+        results = search_engine.query(query)
+        st.subheader("Search Results:")
+        for result in results:
+            st.write(f"Document ID: {result['id']}, Metadata: {result['metadata']}")
+            st.write(f"Text: {result['text']}")
+            st.markdown("---")
+
+if __name__ == "__main__":
+    main()
 

chroma.py

@@ -0,0 +1,39 @@
+from chromadb.utils import embedding_functions
+from chromadb.config import Settings
+import chromadb
+
+class searchengine:
+    def __init__(self, model_name,collection_name):
+        self.sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
+            model_name = model_name
+            )
+        self.chroma_client = chromadb.Client(
+            )
+        self.collection = self.chroma_client.get_or_create_collection(name=collection_name)
+        
+    def add(self, text , metadata,id):
+        self.collection.add(
+            documents = [text],
+            metadatas = [metadata],
+            ids = [id]
+        )
+    
+    def add_list(self, texts:list, metadatas:list, ids:list):
+        self.collection.add(
+            documents = texts,
+            metadatas = metadatas,
+            ids = ids
+        )
+        
+    def query(self, query, number=2):
+        results = self.collection.query(
+            query_texts = [query],
+            n_results = number
+        )
+        return results
+
+    def count(self):
+        return self.collection.count()
+    
+    def peak(self):
+        result = self.collection.peek()