feat(ENDPOINT): :pushpin: Add new summarize endpoint

README.md

@@ -49,6 +49,10 @@ Welcome endpoint that returns a greeting message.
 
 Endpoint to generate text using the language model.
 
+### POST `/summarize`
+
+Endpoint to summarize text using the language model.
+
 **Request parameters:**
 ```json
 {

app.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from functools import partial
 
 from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.checkpoint.memory import MemorySaver
@@ -42,7 +43,7 @@ except Exception as e:
     raise
 
 # Define the function that calls the model
-def call_model(state: MessagesState):
+def call_model(state: MessagesState, system_prompt: str):
     """
     Call the model with the given messages
 
@@ -54,7 +55,7 @@ def call_model(state: MessagesState):
     """
     # Convert LangChain messages to chat format
     messages = [
-        {"role": "system", "content": "You are a friendly Chatbot. Always reply in the language in which the user is writing to you."}
+        {"role": "system", "content": system_prompt}
     ]
     
     for msg in state["messages"]:
@@ -95,12 +96,26 @@ workflow.add_node("model", call_model)
 
 # Add memory
 memory = MemorySaver()
+
+# Define the default system prompt
+DEFAULT_SYSTEM_PROMPT = "You are a friendly Chatbot. Always reply in the language in which the user is writing to you."
+
+# Use partial to create a version of the function with the default system prompt
+workflow.add_node("model", partial(call_model, system_prompt=DEFAULT_SYSTEM_PROMPT))
+
 graph_app = workflow.compile(checkpointer=memory)
 
 # Define the data model for the request
 class QueryRequest(BaseModel):
     query: str
     thread_id: str = "default"
+    system_prompt: str = DEFAULT_SYSTEM_PROMPT
+
+# Define the model for summary requests
+class SummaryRequest(BaseModel):
+    text: str
+    thread_id: str = "default"
+    max_length: int = 200
 
 # Create the FastAPI application
 app = FastAPI(title="LangChain FastAPI", description="API to generate text using LangChain and LangGraph - Máximo Fernández Núñez IriusRisk test challenge")
@@ -119,8 +134,9 @@ async def generate(request: QueryRequest):
     
     Args:
         request: QueryRequest
-        query: str
-        thread_id: str = "default"
+            query: str
+            thread_id: str = "default"
+            system_prompt: str = DEFAULT_SYSTEM_PROMPT
 
     Returns:
         dict: A dictionary containing the generated text and the thread ID
@@ -132,8 +148,12 @@ async def generate(request: QueryRequest):
         # Create the input message
         input_messages = [HumanMessage(content=request.query)]
         
-        # Invoke the graph
-        output = graph_app.invoke({"messages": input_messages}, config)
+        # Invoke the graph with custom system prompt
+        output = graph_app.invoke(
+            {"messages": input_messages}, 
+            config,
+            {"model": {"system_prompt": request.system_prompt}}
+        )
         
         # Get the model response
         response = output["messages"][-1].content
@@ -145,6 +165,47 @@ async def generate(request: QueryRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error generating text: {str(e)}")
 
+@app.post("/summarize")
+async def summarize(request: SummaryRequest):
+    """
+    Endpoint to generate a summary using the language model
+    
+    Args:
+        request: SummaryRequest
+            text: str - The text to summarize
+            thread_id: str = "default"
+            max_length: int = 200 - Maximum summary length
+
+    Returns:
+        dict: A dictionary containing the summary and the thread ID
+    """
+    try:
+        # Configure the thread ID
+        config = {"configurable": {"thread_id": request.thread_id}}
+        
+        # Create a specific system prompt for summarization
+        summary_system_prompt = f"Make a summary of the following text in no more than {request.max_length} words. Keep the most important information and eliminate unnecessary details."
+        
+        # Create the input message
+        input_messages = [HumanMessage(content=request.text)]
+        
+        # Invoke the graph with summarization system prompt
+        output = graph_app.invoke(
+            {"messages": input_messages}, 
+            config,
+            {"model": {"system_prompt": summary_system_prompt}}
+        )
+        
+        # Get the model response
+        response = output["messages"][-1].content
+        
+        return {
+            "summary": response,
+            "thread_id": request.thread_id
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error generating summary: {str(e)}")
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)