Updated memory management and api routing

main.py

@@ -1,25 +1,25 @@
-from fastapi import FastAPI
+from fastapi import FastAPI , UploadFile , File , Form 
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from my_agent.agent import build_graph
 import pandas as pd
-from typing import Optional
+from typing import Optional , List
 from my_agent.utils.initial_interaction import BusinessInteractionChatbot
+import base64
+from PIL import Image
+from io import BytesIO
+import json
 
 app = FastAPI()
 interaction_chatbot = BusinessInteractionChatbot()
 graph = build_graph()
 
 
-class RequestInput(BaseModel):
-    query: list 
-    preferred_topics: Optional[list] = [] 
-
 
 
 
 class UserMessage(BaseModel):
     message: str
-
 details_for_brainstrom = {}
 @app.post("/business-interaction")
 def business_chat(msg: UserMessage):
@@ -32,14 +32,62 @@ def business_chat(msg: UserMessage):
     return {"response": response, "complete": False}
 
 
+
+# class RequestInput(BaseModel):
+#     query: list 
+#     preferred_topics: Optional[list] = [] 
+#     images: Optional[list[str]] = []  # base64-encoded image strings
+
+# @app.post("/brainstrom")
+# def run_graph(input_data: RequestInput):
+#     image_objects = []
+#     for img_b64 in input_data.images:
+#         image_objects.append(process_image(img_b64))  # decode and load images
+
+#     result = graph.invoke({
+#         'topic': input_data.query,
+#         'images': image_objects,
+#         'business_details': details_for_brainstrom
+#     })
+
+#     return {
+#         'final_story': result['final_story'],
+#         'business_details': result['business_details'],
+#     }
+
+
+
+
+# Convert uploaded image to base64 string
+def encode_image_to_base64(uploaded_file: UploadFile) -> str:
+    return base64.b64encode(uploaded_file.file.read()).decode("utf-8")
+
+# Convert base64 string to PIL image (optional for LangGraph processing)
+def process_image(base64_str: str) -> Image.Image:
+    image_data = base64.b64decode(base64_str)
+    return Image.open(BytesIO(image_data))
+
 @app.post("/brainstrom")
-def run_graph(input_data: RequestInput):
-    # business_details = details_for_brainstrom
-    result = graph.invoke({'topic' : input_data.query , 'business_details': details_for_brainstrom})
-    # RequestInput.preferred_topics=result['preferred_topics']
-    return {'final_story': result['final_story'],
-            'business_details':result['business_details'],
-            }
-    
+async def run_graph(
+    query: List[str],  # sent as JSON body
+    preferred_topics: Optional[list] = [],
+    images: Optional[List[UploadFile]] = [],  # ✅ Optional UploadFile list
+    thread_id: Optional[str] = "default-session"
+):
+    # Convert uploaded images to base64
+    image_base64_list = [encode_image_to_base64(img) for img in images]
+
+    # Convert base64 to image objects (if LangGraph expects PIL.Image)
+    image_objects = [process_image(img_b64) for img_b64 in image_base64_list]
 
+    # Invoke LangGraph
+    result = graph.invoke({
+        'topic': query,
+        'images': image_base64_list,
+        'latest_preferred_topics':preferred_topics
+    },
+    config={"configurable": {"thread_id": thread_id}})
 
+    return {
+        'response': result,
+    }

my_agent/agent.py

@@ -1,9 +1,13 @@
 from langgraph.graph import StateGraph, START, END
 from .utils.state import State
-from .utils.nodes import retrieve, generate_story, generate_brainstroming , generate_final_story, route_after_selection, select_preferred_topics
+from .utils.nodes import retrieve, generate_story, generate_brainstroming , generate_final_story, route_after_selection, select_preferred_topics,caption_image
+from langgraph.checkpoint.memory import MemorySaver
 
-def build_graph():
+memory = MemorySaver()
+
+def build_graph_old():
     builder = StateGraph(State)
+    # builder.add_node(caption_image)
     builder.add_node(retrieve)
     builder.add_node(generate_story)
     builder.add_node(generate_brainstroming)
@@ -12,8 +16,13 @@ def build_graph():
 
 
     # Normal edges
+    # builder.add_edge(START, "caption_image")
+
     builder.add_edge(START, "retrieve")
     builder.add_edge("retrieve", "generate_story")
+
+    # builder.add_edge("caption_image", "retrieve")
+    # builder.add_edge("retrieve", "generate_story")
     builder.add_edge("generate_story", "generate_brainstroming")
     builder.add_edge("generate_brainstroming", "select_preferred_topics")
 
@@ -21,4 +30,33 @@ def build_graph():
     builder.add_conditional_edges("select_preferred_topics", route_after_selection,{True:'retrieve',False:'generate_final_story'})
     builder.add_edge("generate_final_story",END)
 
-    return builder.compile()
+    return builder.compile(checkpointer=memory)
+
+
+
+
+def build_graph():
+    builder = StateGraph(State)
+    # builder.add_node(caption_image)
+    builder.add_node(retrieve)
+    builder.add_node(generate_story)
+    builder.add_node(generate_brainstroming)
+    builder.add_node(select_preferred_topics)
+    builder.add_node(generate_final_story)
+
+
+    # Normal edges
+    # builder.add_edge(START, "caption_image")
+
+    builder.add_edge(START, "retrieve")
+    builder.add_edge("retrieve", "generate_story")
+
+    # builder.add_edge("caption_image", "retrieve")
+    # builder.add_edge("retrieve", "generate_story")
+    builder.add_edge("generate_story", "generate_brainstroming")
+
+    # Conditional edge
+    builder.add_edge("generate_brainstroming", END)
+    # builder.add_edge("generate_final_story",END)
+
+    return builder.compile(checkpointer=memory)

my_agent/utils/nodes.py

@@ -5,12 +5,50 @@ from .tools import StoryFormatter, BrainstromTopicFormatter
 from langchain_core.messages import SystemMessage
 from .models_loader import llm , ST
 from .data_loader import load_influencer_data
+from groq import Groq
+import os
+
+
+def caption_image(state: State) -> State:
+    if state.images[-1]!=None:
+      print('Captioning image')
+      client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+
+      chat_completion = client.chat.completions.create(
+          messages=[
+              {
+                  "role": "user",
+                  "content": [
+                      {"type": "text", "text": "What's in this image?"},
+                      {
+                          "type": "image_url",
+                          "image_url": {
+                              "url": f"data:image/jpeg;base64,{state.images[-1]}",
+                          },
+                      },
+                  ],
+              }
+          ],
+          model="meta-llama/llama-4-scout-17b-16e-instruct",
+      )
+      response=chat_completion.choices[0].message.content
+      state.image_captions.append(response)
+      return state
+    
+    else:
+       state.image_captions.append(None)
+       return state
+    
+    # elif state.images[-1]==None:
+    #   state.image_captions.append(None)
+
+    
 
 
 def retrieve(state: State) -> State:
   print('Moving to retrieval process')
   retrievals=[]
-  if len(state.preferred_topics)==0:
+  if len(state.latest_preferred_topics)==0:
       for topic in state.topic:  # Loop through each topic
           embedded_query = ST.encode(topic)  # Embed each topic
           data = load_influencer_data()
@@ -22,7 +60,9 @@ def retrieve(state: State) -> State:
           print('Retrieval process completed......')
       state.retrievals.append(retrievals)
 
-  if len (state.preferred_topics)>0:
+  if len (state.latest_preferred_topics)>0:
+      print('The preferred_topics are:',state.latest_preferred_topics)
+      state.preferred_topics.append(state.latest_preferred_topics)
       for topic in state.preferred_topics[-1]:  # Loop through each topic
           embedded_query = ST.encode(topic)  # Embed each topic
           data = load_influencer_data()
@@ -31,10 +71,11 @@ def retrieve(state: State) -> State:
           # Construct a list of dictionaries for this topic
           result = [{user: story} for user, story in zip(retrieved_examples['username'], retrieved_examples['agentic_story'])]
           retrievals.append(result)
-          print('Retrieval process completed......')
+          print('Retrieval process completed for preferred_topics......')
+      state.latest_preferred_topics=[]
       state.retrievals.append(retrievals)
 
-  print('The retrieval is:\n',state.retrievals )
+#   print('The retrieval is:\n',state.retrievals )
   # return State(messages="Retrieved",topic=state.topic,retrievals=state.retrievals)
   return state
 
@@ -63,6 +104,7 @@ def generate_story(state:State)-> State:
 
   **Final Reminder** You have to strongly focus on these topics while creating the storyline: {state.preferred_topics[-1]}'''
 
+# and {state.image_captions[-1]}
 
   messages = [SystemMessage(content=template)]
   response = llm.bind_tools([StoryFormatter]).invoke(messages)
@@ -145,9 +187,11 @@ def select_preferred_topics(state: State)-> State:
     state.carry_on=True
     return state
 
-
-
-
+def route_after_selection(state:State):
+  if len(state.latest_preferred_topics)==0:
+    return False
+  elif len(state.latest_preferred_topics)>0:
+    return True
 
 def generate_final_story(state:State)-> State:
     if len(state.preferred_topics)>0:
@@ -166,15 +210,15 @@ def generate_final_story(state:State)-> State:
             response = response.content
         else:
           response = "No response"
-        state.final_story=response
+        state.final_story.append(response)
         state.stories.append(response)
         return state
     
-    state.final_story=state.stories[-1]
+    state.final_story.append(state.stories[-1])
+    state.latest_preferred_topics=[]
     return state
 
 
 
-def route_after_selection(state:State):
-    print('The output is:',state.carry_on)
-    return state.carry_on
+
+  

my_agent/utils/state.py

@@ -9,8 +9,10 @@ class State(BaseModel):
     brainstroming_topics: Optional[list] = []
     preferred_topics: Optional[list] = []
     stories : Optional[list]=[]
-    final_story: Optional[str]=None
+    final_story: Optional[list]=[]
     retrievals : Optional[list]=[]
     business_details : Optional[dict]={}
     latest_preferred_topics: Optional[list] = []
+    images: Optional[list[str]] = []  # Base64-encoded strings of images
+    image_captions: Optional[list] = []
     model_config = ConfigDict(arbitrary_types_allowed=True)