feedback in chat added

src/llm.py

@@ -10,6 +10,7 @@ from langgraph.graph import START, MessagesState, StateGraph
 from utils.utils import format_questions_text
 from src.prompts import chat_prompt
 from langchain_openai import ChatOpenAI
+from src.models import Phase2Generation
 @dataclass
 class Message:
     role: str  # 'human' or 'ai'
@@ -37,10 +38,10 @@ def initialize_workflow(model) -> StateGraph:
     return workflow.compile(checkpointer=memory)
 
 
-def setup_prompt_template(theme: int, resume: str) -> ChatPromptTemplate:
+def setup_prompt_template(theme: int, resume: str,full_history=None,form_response=None,generate_theme="NO") -> ChatPromptTemplate:
     """Set up the prompt template"""
     return ChatPromptTemplate.from_messages([
-        ("system", chat_prompt(theme, resume)),
+        ("system", chat_prompt(theme, resume,full_history,form_response,generate_theme)),
         MessagesPlaceholder(variable_name="messages")
     ])
 
@@ -112,7 +113,7 @@ def convert_to_langchain_messages(messages: List[Message]) -> List[HumanMessage
     return converted_messages
 
 
-def ai_chat(query: str, conversation_id: str, theme_id: int, resume: str) -> str:
+def ai_chat(query: str, conversation_id: str, theme_id: int, resume: str,full_history=None,form_response=None,generate_theme="NO") -> str:
     """Main chat function that processes queries and manages conversation"""
     storage_path = Path("conversations.json")
     
@@ -121,12 +122,12 @@ def ai_chat(query: str, conversation_id: str, theme_id: int, resume: str) -> str
         language: str
     # Initialize model and workflow
     model = ChatOpenAI(model=MODEL)
+    if generate_theme == "YES":
+        model = model.with_structured_output(Phase2Generation)
     workflow = StateGraph(state_schema=State)
     
-    
-    
     def call_model(state: State):
-        prompt_template = setup_prompt_template(theme_id, resume)
+        prompt_template = setup_prompt_template(theme_id, resume,full_history,form_response,generate_theme)
         prompt = prompt_template.invoke({
             "messages": state["messages"], 
             "language": state["language"]
@@ -137,6 +138,7 @@ def ai_chat(query: str, conversation_id: str, theme_id: int, resume: str) -> str
     workflow.add_edge(START, "model")
     workflow.add_node("model", call_model)
     
+
     memory = MemorySaver()
     app = workflow.compile(checkpointer=memory)
 
@@ -162,12 +164,21 @@ def ai_chat(query: str, conversation_id: str, theme_id: int, resume: str) -> str
             config
         )
 
+    if generate_theme == "YES":
+        structured_message = output["messages"][0]
+        output = structured_message.json(by_alias=True)  # This returns a JSON string.
+        print(f"Output: {output}")
+        add_message(storage_path, conversation_id, "ai", output)
+    else:
+        output = output["messages"][-1].content
+        add_message(storage_path, conversation_id, "ai", output)
+
     # Store messages
     if query:
         add_message(storage_path, conversation_id, "human", query)
-    add_message(storage_path, conversation_id, "ai", output["messages"][-1].content)
-
-    return output["messages"][-1].content
+    
+    print(output)
+    return output
 
 
 # Example usage: