feature(#35): implement train logic.

Brain/src/common/utils.py

@@ -30,8 +30,8 @@ DEFAULT_GPT_MODEL = "gpt-4"
 AGENT_NAME = "RisingBrain Assistant"
 
 # indexes of relatedness of embedding
-COMMAND_SMS_INDEXS = [4, 5]
-COMMAND_BROWSER_OPEN = [10]
+COMMAND_SMS_INDEXES = ["pWDrks5DO1bEPLlUtQ1f", "LEpAhmFi8tAOQUE7LHZZ"]  # 4, 5
+COMMAND_BROWSER_OPEN = ["taVNeDINonUqJWXBlESU"]  # 10
 
 # Twilio
 ACCOUNT_SID = os.getenv("TWILIO_ACCOUNT_SID")
@@ -86,8 +86,10 @@ def parseJsonFromCompletion(data: str) -> json:
     if index == len(result) - 3:
         result = result[:index] + replacement + result[index + len(substring):]
     # fmt: on
-    result = json.loads(result.replace("':", '":'))
-    return result
+    try:
+        return json.loads(result)
+    except Exception as e:
+        return result
 
 
 def parseUrlFromStr(text: str) -> str:

Brain/src/model/requests/request_model.py

@@ -130,6 +130,14 @@ class TrainContacts(BasicReq):
     contacts: list[ContactReq]
 
 
+"""endpoint: /document"""
+
+
+class Document(BasicReq):
+    document_id: str
+    page_content: str
+
+
 """endpoint /browser/item"""
 
 
@@ -140,3 +148,16 @@ class BrowserItem(BasicReq):
 
     items: list[ItemReq]
     prompt: str
+
+
+"""endpoint /train"""
+
+
+class Train(BasicReq):
+    class TrainData(BaseModel):
+        page_content: str
+        timestamp: float
+
+    id: str
+    data: TrainData
+    status: str

Brain/src/model/train_model.py

@@ -0,0 +1,24 @@
+"""train model:
+{
+    "id": "String",
+    "data": [{"page_content": "String", "timestamp": 0}],
+    "status": "created | updated | deleted",
+}"""
+
+from Brain.src.model.requests.request_model import Train
+
+
+class TrainModel:
+    def __init__(self, train_data: Train):
+        self.id = train_data.id
+        self.data = train_data.data
+        self.status = TrainStatus.UPDATED
+
+
+"""train status: created | updated | deleted"""
+
+
+class TrainStatus:
+    CREATED = "created"
+    UPDATED = "updated"
+    DELETED = "deleted"

Brain/src/rising_plugin/guardrails-config/actions/actions.py

@@ -13,20 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import json
-import numpy as np
 
 from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.vectorstores import utils
-from langchain.document_loaders.csv_loader import CSVLoader
+from Brain.src.service.train_service import TrainService
 from langchain.docstore.document import Document
 
 from Brain.src.common.brain_exception import BrainException
 from Brain.src.common.utils import (
     OPENAI_API_KEY,
-    COMMAND_SMS_INDEXS,
+    COMMAND_SMS_INDEXES,
     COMMAND_BROWSER_OPEN,
+    PINECONE_INDEX_NAME,
 )
 from Brain.src.rising_plugin.image_embedding import (
     query_image_text,
@@ -44,20 +42,20 @@ from Brain.src.rising_plugin.llm.llms import (
     FALCON_7B,
 )
 
-"""
-query is json string with below format
-{
-    "query": string,
-    "model": string,
-    "uuid": string,
-    "image_search": bool,
-}
-"""
+from Brain.src.rising_plugin.pinecone_engine import (
+    get_pinecone_index_namespace,
+    init_pinecone,
+)
+
+
+def get_pinecone_index_train_namespace() -> str:
+    return get_pinecone_index_namespace(f"trains")
 
 
 @action()
 async def general_question(query):
     """step 0: convert string to json"""
+    index = init_pinecone(PINECONE_INDEX_NAME)
     try:
         json_query = json.loads(query)
     except Exception as ex:
@@ -69,25 +67,26 @@ async def general_question(query):
     image_search = json_query["image_search"]
 
     """step 1: handle with gpt-4"""
-    file_path = os.path.dirname(os.path.abspath(__file__))
 
-    with open(f"{file_path}/phone.json", "r") as infile:
-        data = json.load(infile)
     embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
 
     query_result = embeddings.embed_query(query)
-    doc_list = utils.maximal_marginal_relevance(np.array(query_result), data, k=1)
-    loader = CSVLoader(file_path=f"{file_path}/phone.csv", encoding="utf8")
-    csv_text = loader.load()
+    relatedness_data = index.query(
+        vector=query_result,
+        top_k=3,
+        include_values=False,
+        namespace=get_pinecone_index_train_namespace(),
+    )
+    documentId = ''
+    if len(relatedness_data["matches"]) > 0:
+        documentId = relatedness_data["matches"][0]["id"]
 
     docs = []
-
-    for res in doc_list:
-        docs.append(
-            Document(
-                page_content=csv_text[res].page_content, metadata=csv_text[res].metadata
-            )
-        )
+    train_service = TrainService()
+    documents = train_service.read_all_documents()
+    for document in documents:
+        if document["document_id"] == documentId:
+            docs.append(Document(page_content=document["page_content"], metadata=""))
 
     chain_data = get_llm_chain(model=model).run(input_documents=docs, question=query)
     # test
@@ -115,8 +114,8 @@ async def general_question(query):
         return str(result)
     except ValueError as e:
         # Check sms and browser query
-        if doc_list[0] in COMMAND_SMS_INDEXS:
+        if documentId in COMMAND_SMS_INDEXES:
             return str({"program": "sms", "content": chain_data})
-        elif doc_list[0] in COMMAND_BROWSER_OPEN:
+        elif documentId in COMMAND_BROWSER_OPEN:
             return str({"program": "browser", "content": "https://google.com"})
         return str({"program": "message", "content": falcon_llm.query(question=query)})

Brain/src/router/api.py

@@ -24,7 +24,6 @@ from Brain.src.rising_plugin.risingplugin import (
     handle_chat_completion,
 )
 from Brain.src.firebase.cloudmessage import send_message, get_tokens
-from Brain.src.rising_plugin.csv_embed import csv_embed
 from Brain.src.rising_plugin.image_embedding import embed_image_text, query_image_text
 
 from Brain.src.logs import logger
@@ -162,16 +161,6 @@ def construct_blueprint_api() -> APIRouter:
             result={"program": "image", "content": image_response},
         )
 
-    """@generator.response(
-        status_code=200, schema={"message": "message", "result": "test_result"}
-    )"""
-
-    @router.get("/training")
-    def csv_training():
-        csv_embed()
-
-        return assembler.to_response(200, "trained successfully", "")
-
     """@generator.request_body(
         {
             "token": "test_token",

Brain/src/router/train_router.py

@@ -0,0 +1,104 @@
+from fastapi import APIRouter
+
+from Brain.src.common.assembler import Assembler
+from Brain.src.model.requests.request_model import (
+    Document,
+)
+from Brain.src.service.train_service import TrainService
+
+router = APIRouter()
+
+
+def construct_blueprint_train_api() -> APIRouter:
+    # Assembler
+    assembler = Assembler()
+
+    # Services
+    train_service = TrainService()
+
+    """@generator.response(
+        status_code=200, schema={"message": "message", "result": "test_result"}
+    )"""
+
+    @router.get("")
+    def read_all_documents():
+        try:
+            result = train_service.read_all_documents()
+        except Exception as e:
+            return assembler.to_response(400, "failed to get all documents", "")
+        return assembler.to_response(200, "Get all documents list successfully", result)
+
+    """@generator.response( status_code=200, schema={"message": "message", "result": {"document_id": "document_id", 
+    "page_content":"page_content"}} )"""
+
+    @router.get("/{document_id}")
+    def read_one_document(document_id: str):
+        try:
+            result = train_service.read_one_document(document_id)
+        except Exception as e:
+            return assembler.to_response(400, "fail to get one document", "")
+        return assembler.to_response(200, "Get one document successfully", result)
+
+    """@generator.request_body(
+        {
+            "token": "test_token",
+            "uuid": "test_uuid",
+            "page_content": "string",            
+        }
+    )
+    @generator.response( status_code=200, schema={"message": "message", "result": {"document_id": "document_id", 
+    "page_content":"page_content"}} )"""
+
+    @router.post("")
+    def create_document_train(data: Document):
+        try:
+            result = train_service.create_one_document(data.page_content)
+        except Exception as e:
+            return assembler.to_response(400, "failed to create one document", "")
+        return assembler.to_response(
+            200, "created one document and trained it successfully", result
+        )
+
+    """@generator.request_body(
+        {
+            "token": "test_token",
+            "uuid": "test_uuid",
+            "document_id": "string",
+            "page_content": "string",
+        }
+    )
+    @generator.response( status_code=200, schema={"message": "message", "result": {"document_id": "document_id", 
+    "page_content":"page_content"}} )"""
+
+    @router.put("")
+    def update_one_document(data: Document):
+        try:
+            result = train_service.update_one_document(
+                data.document_id, data.page_content
+            )
+        except Exception as e:
+            return assembler.to_response(400, "fail to update one document", "")
+        return assembler.to_response(
+            200, "updated one document and trained it successfully", result
+        )
+
+    """@generator.request_body(
+        {
+            "token": "test_token",
+            "uuid": "test_uuid",
+            "document_id": "string",            
+        }
+    )
+    @generator.response( status_code=200, schema={"message": "message", "result": {"document_id": "document_id"}} )"""
+
+    @router.delete("/{document_id}")
+    def delete_one_document(document_id: str):
+        try:
+            result = train_service.delete_one_document(document_id)
+        except Exception as e:
+            return assembler.to_response(400, "fail to delete one train", "")
+        return assembler.to_response(
+            200, "deleted one document and train data successfully", result
+        )
+
+    return router

Brain/src/service/train_service.py

@@ -0,0 +1,117 @@
+"""service to manage trains"""
+from typing import List, Any
+
+from Brain.src.rising_plugin.csv_embed import get_embed
+from Brain.src.rising_plugin.pinecone_engine import (
+    get_pinecone_index_namespace,
+    update_pinecone,
+    init_pinecone,
+    delete_pinecone,
+    add_pinecone,
+    delete_all_pinecone,
+)
+
+from firebase_admin import firestore
+import datetime
+
+
+def to_json(page_content: str):
+    return {
+        "page_content": page_content,
+        "timestamp": datetime.datetime.now().timestamp(),
+    }
+
+
+class TrainService:
+    """train (getting embedding) and update pinecone with embeddings by train_id
+    train datatype:
+    key: id
+    values: {id},{data}, {status}"""
+
+    def __init__(self):
+        self.db = firestore.client()
+        self.documents_ref = self.db.collection("documents")
+
+    """read all documents from firestore"""
+
+    def read_all_documents(self):
+        query = self.documents_ref.order_by("timestamp")
+        docs = query.stream()
+        result = []
+        for item in docs:
+            item_data = item.to_dict()
+            result.append(
+                {"document_id": item.id, "page_content": item_data["page_content"]}
+            )
+        return result
+
+    """read one document from firestore"""
+
+    def read_one_document(self, document_id: str):
+        doc = self.documents_ref.document(document_id).get()
+        if doc.exists:
+            return {
+                "document_id": document_id,
+                "page_content": doc.to_dict()["page_content"],
+            }
+        else:
+            return None
+
+    """create a new document and train it"""
+
+    def create_one_document(self, page_content: str):
+        # Auto-generate document ID
+        auto_generated_doc_ref = self.documents_ref.document()
+        auto_generated_doc_ref.set(to_json(page_content))
+        auto_generated_document_id = auto_generated_doc_ref.id
+        self.train_one_document(auto_generated_document_id, page_content)
+        return {"document_id": auto_generated_document_id, "page_content": page_content}
+
+    """update a document by using id and train it"""
+
+    def update_one_document(self, document_id: str, page_content: str):
+        self.documents_ref.document(document_id).update(to_json(page_content))
+        self.train_one_document(document_id, page_content)
+        return {"document_id": document_id, "page_content": page_content}
+
+    """delete a document by using document_id"""
+
+    def delete_one_document(self, document_id: str):
+        self.documents_ref.document(document_id).delete()
+        self.delete_one_pinecone(document_id)
+        return {"document_id": document_id}
+
+    def train_all_documents(self) -> str:
+        documents = self.read_all_documents()
+        result = list()
+        pinecone_namespace = self.get_pinecone_index_namespace()
+        for item in documents:
+            query_result = get_embed(item["page_content"])
+            result.append(query_result)
+            key = item["document_id"]
+            value = f'{item["page_content"]}, {query_result}'
+            # get vectoring data(embedding data)
+            vectoring_values = get_embed(value)
+            add_pinecone(namespace=pinecone_namespace, key=key, value=vectoring_values)
+
+        return "trained all documents successfully"
+
+    def train_one_document(self, document_id: str, page_content: str) -> None:
+        pinecone_namespace = self.get_pinecone_index_namespace()
+        result = list()
+        query_result = get_embed(page_content)
+        result.append(query_result)
+        key = document_id
+        value = f"{page_content}, {query_result}"
+        # get vectoring data(embedding data)
+        vectoring_values = get_embed(value)
+        add_pinecone(namespace=pinecone_namespace, key=key, value=vectoring_values)
+
+    def delete_all(self) -> Any:
+        return delete_all_pinecone(self.get_pinecone_index_namespace())
+
+    def delete_one_pinecone(self, document_id: str) -> Any:
+        return delete_pinecone(self.get_pinecone_index_namespace(), document_id)
+
+    def get_pinecone_index_namespace(self) -> str:
+        return get_pinecone_index_namespace(f"trains")

app.py

@@ -3,6 +3,7 @@ from fastapi import Depends, FastAPI
 import uvicorn
 
 from Brain.src.router.browser_router import construct_blueprint_browser_api
+from Brain.src.router.train_router import construct_blueprint_train_api
 
 initialize_app()
 
@@ -14,6 +15,10 @@ app.include_router(
     construct_blueprint_browser_api(), prefix="/browser", tags=["ai_browser"]
 )
 
+app.include_router(
+    construct_blueprint_train_api(), prefix="/train", tags=["ai_train"]
+)
+
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)