refactor(app.py): :sparkles: Improve response handling and streamline configuration for model invocation

app.py

@@ -78,13 +78,15 @@ def call_model(state: MessagesState, system_prompt: str):
         pad_token_id=tokenizer.eos_token_id
     )
     
-    # Decode and clean the response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Extract only the assistant's response (after the last user message)
-    response = response.split("Assistant:")[-1].strip()
+    # Get just the new tokens (excluding the input prompt tokens)
+    input_length = inputs.shape[1]
+    generated_tokens = outputs[0][input_length:]
+    
+    # Decode only the new tokens to get just the assistant's response
+    assistant_response = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
     
     # Convert the response to LangChain format
-    ai_message = AIMessage(content=response)
+    ai_message = AIMessage(content=assistant_response)
     return {"messages": state["messages"] + [ai_message]}
 
 # Define the graph
@@ -92,7 +94,6 @@ workflow = StateGraph(state_schema=MessagesState)
 
 # Define the node in the graph
 workflow.add_edge(START, "model")
-workflow.add_node("model", call_model)
 
 # Add memory
 memory = MemorySaver()
@@ -149,10 +150,16 @@ async def generate(request: QueryRequest):
         input_messages = [HumanMessage(content=request.query)]
         
         # Invoke the graph with custom system prompt
+        # Combine config parameters into a single dictionary
+        combined_config = {
+            **config,
+            "model": {"system_prompt": request.system_prompt}
+        }
+        
+        # Invoke the graph with proper argument count
         output = graph_app.invoke(
             {"messages": input_messages}, 
-            config,
-            {"model": {"system_prompt": request.system_prompt}}
+            combined_config
         )
         
         # Get the model response
@@ -189,11 +196,16 @@ async def summarize(request: SummaryRequest):
         # Create the input message
         input_messages = [HumanMessage(content=request.text)]
         
-        # Invoke the graph with summarization system prompt
+        # Combine config parameters into a single dictionary
+        combined_config = {
+            **config,
+            "model": {"system_prompt": summary_system_prompt}
+        }
+        
+        # Invoke the graph with proper argument count
         output = graph_app.invoke(
             {"messages": input_messages}, 
-            config,
-            {"model": {"system_prompt": summary_system_prompt}}
+            combined_config
         )
         
         # Get the model response