feat: Added langflow components and prompt API

Author: rohanpooniwala

.vscode/settings.json

@@ -0,0 +1,11 @@
+{
+    "python.linting.enabled": true,
+    "python.formatting.blackArgs": [
+        "--line-length",
+        "140"
+    ],
+    "python.formatting.provider": "black",
+    "editor.formatOnSave": true,
+    "python.linting.mypyEnabled": false,
+    "python.analysis.typeCheckingMode": "basic"
+}

requirements.txt

@@ -4,3 +4,5 @@ openai
 sqlalchemy
 pyjwt[crypto]
 ipython
+langflow
+black

server/api/chatbot.py

@@ -8,10 +8,12 @@
 
 chatbot_router = APIRouter(prefix="/chatbot", tags=["chatbot"])
 
+
 class ChatBotModel(BaseModel):
     name: str
     dag: dict
 
+
 @chatbot_router.post("/", status_code=200)
 def create_chatbot(inputs: ChatBotModel, token: Annotated[str, Header()], db: Session = Depends(database.db_session)):
     username = get_user_from_jwt(token)
@@ -26,18 +28,19 @@ def create_chatbot(inputs: ChatBotModel, token: Annotated[str, Header()], db: Se
         response = {"msg": "failed"}
     return response
 
+
 # @chatbot_router.get("/", status_code=200)
 
 # @chatbot_router.put("/", status_code=200)
 # def update_chatbot(inputs: ChatBotModel, db: Session = Depends(database.db_session)):
-    # user: User = db.query(User).filter((User.username == inputs.username) & (User.password == inputs.old_password)).first()
-    # if user is not None:
-    #     user.password = inputs.new_password
-    #     db.commit()
-    #     response = {"msg": "success"}
-    # else:
-    #     response = {"msg": "failed"}
-    # return response
+# user: User = db.query(User).filter((User.username == inputs.username) & (User.password == inputs.old_password)).first()
+# if user is not None:
+#     user.password = inputs.new_password
+#     db.commit()
+#     response = {"msg": "success"}
+# else:
+#     response = {"msg": "failed"}
+# return response
 
 # @chatbot_router.delete("/", status_code=200)
 # def update_chatbot(inputs: ChatBotModel, db: Session = Depends(database.db_session)):

server/api/langflow.py

@@ -0,0 +1,21 @@
+import logging
+from typing import Any, Dict
+
+from fastapi import APIRouter, HTTPException
+
+from langflow.interface.run import process_graph
+from langflow.interface.types import build_langchain_types_dict
+
+# build router
+router = APIRouter(prefix="/flow", tags=["flow"])
+# add docs to router
+router.__doc__ = """
+# Flow API
+"""
+
+logger = logging.getLogger(__name__)
+
+
+@router.get("/components")
+def get_all():
+    return build_langchain_types_dict()

server/api/prompts.py

@@ -0,0 +1,127 @@
+import logging
+from typing import Any, Dict
+from sqlalchemy.orm import Session
+from database import db_session
+import contextlib
+import io
+
+from fastapi import APIRouter, HTTPException, Depends
+
+from langflow.interface.run import load_langchain_object, save_cache, fix_memory_inputs
+
+from database_utils.chatbot import get_chatbot
+from schemas.prompt_schema import PromptSchema
+
+# build router
+router = APIRouter(tags=["prompts"])
+# add docs to router
+router.__doc__ = """
+# Prompts API
+"""
+
+logger = logging.getLogger(__name__)
+
+
+def format_intermediate_steps(intermediate_steps):
+    formatted_chain = []
+    for step in intermediate_steps:
+        action = step[0]
+        observation = step[1]
+
+        formatted_chain.append(
+            {
+                "action": action.tool,
+                "action_input": action.tool_input,
+                "observation": observation,
+            }
+        )
+    return formatted_chain
+
+
+def get_result_and_thought_using_graph(langchain_object, message: str):
+    """Get result and thought from extracted json"""
+    try:
+        if hasattr(langchain_object, "verbose"):
+            langchain_object.verbose = True
+        chat_input = None
+        memory_key = ""
+        if hasattr(langchain_object, "memory") and langchain_object.memory is not None:
+            memory_key = langchain_object.memory.memory_key
+
+        for key in langchain_object.input_keys:
+            if key not in [memory_key, "chat_history"]:
+                chat_input = {key: message}
+
+        if hasattr(langchain_object, "return_intermediate_steps"):
+            # https://github.com/hwchase17/langchain/issues/2068
+            # Deactivating until we have a frontend solution
+            # to display intermediate steps
+            langchain_object.return_intermediate_steps = True
+
+        fix_memory_inputs(langchain_object)
+
+        with io.StringIO() as output_buffer, contextlib.redirect_stdout(output_buffer):
+            try:
+                output = langchain_object(chat_input)
+            except ValueError as exc:
+                # make the error message more informative
+                logger.debug(f"Error: {str(exc)}")
+                output = langchain_object.run(chat_input)
+
+            intermediate_steps = output.get("intermediate_steps", []) if isinstance(output, dict) else []
+
+            result = output.get(langchain_object.output_keys[0]) if isinstance(output, dict) else output
+            if intermediate_steps:
+                thought = format_intermediate_steps(intermediate_steps)
+            else:
+                thought = {"steps": output_buffer.getvalue()}
+
+    except Exception as exc:
+        raise ValueError(f"Error: {str(exc)}") from exc
+    return result, thought
+
+
+def process_graph(message, chat_history, data_graph):
+    """
+    Process graph by extracting input variables and replacing ZeroShotPrompt
+    with PromptTemplate,then run the graph and return the result and thought.
+    """
+    # Load langchain object
+    logger.debug("Loading langchain object")
+    is_first_message = len(chat_history) == 0
+    computed_hash, langchain_object = load_langchain_object(data_graph, is_first_message)
+    logger.debug("Loaded langchain object")
+
+    if langchain_object is None:
+        # Raise user facing error
+        raise ValueError("There was an error loading the langchain_object. Please, check all the nodes and try again.")
+
+    # Generate result and thought
+    logger.debug("Generating result and thought")
+    result, thought = get_result_and_thought_using_graph(langchain_object, message)
+    logger.debug("Generated result and thought")
+
+    # Save langchain_object to cache
+    # We have to save it here because if the
+    # memory is updated we need to keep the new values
+    logger.debug("Saving langchain object to cache")
+    save_cache(computed_hash, langchain_object, is_first_message)
+    logger.debug("Saved langchain object to cache")
+    return {"result": str(result), "thought": thought}
+
+
+@router.post("/chatbot/{chatbot_id}/prompt")
+def get_prompt(chatbot_id: int, prompt: PromptSchema, db: Session = Depends(db_session)):
+    try:
+        chatbot = get_chatbot(db, chatbot_id)
+
+        # Process graph
+        logger.debug("Processing graph")
+        result = process_graph(prompt.new_message, prompt.chat_history, chatbot.dag)
+
+        logger.debug("Processed graph")
+        return result
+
+    except Exception as e:
+        logger.exception(e)
+        raise HTTPException(status_code=500, detail=str(e)) from e

server/app.py

@@ -14,6 +14,8 @@
 from api.metrics import metrics_router
 from api.feedback import feedback_router
 from api.intermediate_steps import intermediate_steps_router
+from api.langflow import router as langflow_router
+from api.prompts import router as prompts_router
 
 c.Base.metadata.create_all(bind=engine)
 
@@ -34,7 +36,8 @@
 app.include_router(intermediate_steps_router)
 app.include_router(chatbot_router)
 app.include_router(auth_router)
-
+app.include_router(langflow_router)
+app.include_router(prompts_router)
 ####################################################
 ################ APIs ##############################
 ####################################################

server/database.py

@@ -60,8 +60,8 @@ class Prompt(Base):
     gpt_rating = Column(Text, nullable=True)
     user_rating = Column(Enum(PromptRating), nullable=True)
     chatbot_user_rating = input_prompt = Column(Enum(PromptRating), nullable=True)
-    response = Column(Text, nullable=False)
-    time_taken = Column(Float, nullable=False)
+    response = Column(Text, nullable=True)
+    time_taken = Column(Float, nullable=True)
     created_at = Column(DateTime, nullable=False)
     session_id = Column(String(80), nullable=False)
     meta = Column(JSON)

server/database_utils/__init__.py

@@ -0,0 +1,12 @@
+from sqlalchemy.orm import Session
+
+from database import ChatBot
+
+
+def get_chatbot(db: Session, chatbot_id: int) -> ChatBot:
+    row = db.query(ChatBot).filter(ChatBot.id == chatbot_id).first()
+
+    if row is None:
+        raise ValueError(f"Chatbot with id {chatbot_id} does not exist")
+
+    return row

server/database_utils/chatbot.py

@@ -0,0 +1,23 @@
+from datetime import datetime
+from typing import List
+
+from sqlalchemy.orm import Session
+
+from database import ChatBot, Prompt, IntermediateStep
+
+
+def get_prompts(db: Session, chatbot_id: int) -> List[Prompt]:
+    row = db.query(Prompt).filter(Prompt.chatbot_id == chatbot_id).all()
+    return row
+
+
+def create_prompt(db: Session, input_prompt: str, session_id=str) -> Prompt:
+    db_prompt = Prompt(
+        input_prompt=input_prompt,
+        created_at=datetime.now(),
+        session_id=session_id,
+    )
+    db.add(db_prompt)
+    db.commit()
+    db.refresh(db_prompt)
+    return db_prompt