Upload 3 files

receipt_gen_agent.py

@@ -0,0 +1,260 @@
+from langchain_core.output_parsers import JsonOutputParser
+from langchain_core.prompts import PromptTemplate
+from dotenv import load_dotenv
+import os
+from typing import List
+from typing_extensions import TypedDict
+from langchain_core.messages import HumanMessage
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.output_parsers import RetryOutputParser
+from langgraph.graph import StateGraph, START, END
+import base64
+from IPython.display import Image as img, display
+from langchain_core.runnables.graph import MermaidDrawMethod
+from langgraph.checkpoint.memory import MemorySaver
+import json
+from pydantic import BaseModel, Field
+from io import BytesIO
+load_dotenv()
+GEMINI_API_KEY=os.getenv('google_api_key')
+
+
+GEMINI_MODEL='gemini-2.0-flash'
+llm = ChatGoogleGenerativeAI(google_api_key=GEMINI_API_KEY, model=GEMINI_MODEL, temperature=0.3)
+
+from os import listdir
+from os.path import isfile, join
+
+
+class State(TypedDict):
+    prompt: str
+    image_number: int
+    image_data: json
+    image_byte: str
+    eval: dict
+    n_retries:int
+    image_name: str
+    image_data_list: list
+
+
+def generate_data_node(state:State):
+    class Items(BaseModel):
+        name: str = Field(description='the name of the item')
+        price : float = Field(description='the price of the item')
+        quantity: int = Field(description='the quantity of the item')
+
+    class Form(BaseModel):
+        loc_name: str = Field(description='the name of the location if no name put empty str')
+        address: str = Field(description='the address of the location if no location put empty str')
+        date: str = Field(description='the date if no date put empty str')
+        time: str = Field(description='the time if no time put empty str')
+        items: List[Items] = Field(description= 'list of the items if no items put empty list')
+        subtotal: float = Field(description= 'the subtotal if no subtotal put 0')
+        tax: float = Field(description='the tax, if no tax put 0')
+        total: float = Field(description='the total amount if no total amount put 0')
+
+        
+    parser=JsonOutputParser(pydantic_object=Form)
+    instruction=parser.get_format_instructions()
+    message = HumanMessage(
+    content=[
+        {"type": "text", "text": f"{state.get('prompt')}"+'\n\n'+ instruction},
+        {
+            "type": "image_url",
+            "image_url": {"url": f"data:image/jpeg;base64,{state.get('image_byte')}"},
+        },
+    ],
+)
+    response=llm.invoke([message])
+    try:
+        response=parser.parse(response.content)
+        return {'image_data':response}
+    except:
+        prompt = PromptTemplate(
+        template="Answer the user query.\n{format_instructions}\n{query}\n",
+        input_variables=["query"],
+        partial_variables={"format_instructions": parser.get_format_instructions()},
+        )
+        retry_parser = RetryOutputParser.from_llm(parser=parser, llm=llm)
+        prompt_value=prompt.format_prompt(query=f'{state.get('prompt')}')
+        response=retry_parser.parse_with_prompt(response.content, prompt_value)     
+        return {'image_data':response}
+    
+def evaluate_node(state:State):
+
+
+    class Decision(BaseModel):
+        decision: str = Field(description='good or modify if changes have to be made')
+        comment: str = Field(description='the changes to make')
+        
+    parser=JsonOutputParser(pydantic_object=Decision)
+    prompt = PromptTemplate(
+    template="Answer the user query.\n{format_instructions}\n{query}\n",
+    input_variables=["query"],
+    partial_variables={"format_instructions": parser.get_format_instructions()},
+    )
+    data=state.get('image_data')
+    query=f" is the {data} correct and makes sense tell the llm what to change, ignore missing data, don't make it up, no explanation or decription needed"
+    chain = prompt | llm 
+    response=chain.invoke({'query':query}) 
+    try:
+        response=parser.parse(response.content)     
+    except:
+
+        retry_parser = RetryOutputParser.from_llm(parser=parser, llm=llm)
+
+        prompt_value = prompt.format_prompt(query=query)
+        response=retry_parser.parse_with_prompt(response.content, prompt_value)       
+    return {'eval': response}
+
+
+def data_editor_node(state:State):
+    class Items(BaseModel):
+        name: str = Field(description='the name of the item')
+        price : float = Field(description='the price of the item')
+        quantity: int = Field(description='the quantity of the item')
+
+    class Form(BaseModel):
+        loc_name: str = Field(description='the name of the location if no name put empty str')
+        address: str = Field(description='the address of the location if no location put empty str')
+        date: str = Field(description='the date if no date put empty str')
+        time: str = Field(description='the time if no time put empty str')
+        items: List[Items] = Field(description= 'list of the items if no items put empty list')
+        subtotal: float = Field(description= 'the subtotal if no subtotal put 0')
+        tax: float = Field(description='the tax, if no tax put 0')
+        total: float = Field(description='the total amount if no total amount put 0')
+
+        
+    parser=JsonOutputParser(pydantic_object=Form)
+    prompt = PromptTemplate(
+    template="Answer the user query.\n{format_instructions}\n{query}\n",
+    input_variables=["query"],
+    partial_variables={"format_instructions": parser.get_format_instructions()},
+    )
+
+
+    data=state.get('image_data')
+    query=f"modify this dict: {data} based on these comments {state.get('eval').get('comment')}, return a json"
+    chain = prompt | llm 
+    response=chain.invoke({'query':query}) 
+    try:
+        response=parser.parse(response.content)     
+    except:
+
+        retry_parser = RetryOutputParser.from_llm(parser=parser, llm=llm)
+
+        prompt_value = prompt.format_prompt(query=query)
+        response=retry_parser.parse_with_prompt(response.content, prompt_value)         
+    return {'image_data': response,
+            'n_retries':state.get('n_retries')+1}
+
+
+def should_continue(state:State)-> str:
+    """
+        Determine whether the research process should continue based on the current state.
+
+        Args:
+            state: The current state of the agent.
+
+        Returns:
+            str: The next state to transition to ("to_add_data", "to_prompt_editor").
+        """
+    eval=state.get('eval').get('decision')
+    if eval =='good':
+        return 'to_add_data'
+
+    elif eval =='modify' and state.get('n_retries')<2:
+        return 'to_data_editor'
+    else:
+        return 'to_add_data'
+    
+
+def add_data_node(state:State):
+        img_number=state.get('image_number')
+        return {
+                'n_retries':0,
+                
+                'image_name':f'{img_number}_new_receipt.jpg'}
+
+class receipt_agent:
+    def __init__(self):
+        self.agent=self._setup()
+    def _setup(self):
+   
+        agent_builder=StateGraph(State)
+        agent_builder.add_node('generate_data',generate_data_node)
+        agent_builder.add_node('evaluate',evaluate_node)
+        agent_builder.add_node('add_data',add_data_node)
+        agent_builder.add_node('data_editor',data_editor_node)
+
+        agent_builder.add_edge(START,'generate_data')
+        agent_builder.add_edge('generate_data','evaluate')
+        # agent_builder.add_edge('evaluate',END)
+        agent_builder.add_conditional_edges('evaluate', should_continue, {'to_data_editor':'data_editor', 'to_add_data':'add_data'},)
+        agent_builder.add_edge('data_editor','evaluate')
+        agent_builder.add_edge('add_data', END)
+
+
+        checkpointer=MemorySaver()
+
+        agent=agent_builder.compile(checkpointer=checkpointer)
+        return agent
+        
+
+    def display_graph(self):
+        return display(
+                        img(
+                                self.agent.get_graph().draw_mermaid_png(
+                                    draw_method=MermaidDrawMethod.API,
+                                )
+                            )
+                        )
+    
+    def get_state(self, state_val:str):
+        config = {"configurable": {"thread_id": "1"}}
+        return self.agent.get_state(config).values[state_val]
+    
+    def receipt_gen(self,image):
+        config = {"configurable": {"thread_id": "1"}}
+        buffered=BytesIO()
+
+        image.save(buffered, format='JPEG')
+        image_data = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+        data_list = [f for f in listdir('new_receipt_data') if isfile(join('new_receipt_data', f))]
+        if not data_list:
+            data_list=[]
+        else:
+            with open(f'new_receipt_data/{data_list[0]}', 'r') as openfile:
+            # Reading from json file
+                data_list = json.load(openfile)
+            
+        response=self.agent.invoke({'prompt':'analyse this receipt and list the items, return a json',
+                                'n_retries':0,
+                                'image_number':len(data_list),
+                                'image_byte': image_data,
+                                'image_data_list':data_list}, config)
+        
+        image_data=response.get('image_data')
+        return image_data
+    
+    def update_state(self, values:dict):
+        config = {"configurable": {"thread_id": "1"}}
+        return self.agent.update_state(config,values=values)
+    
+    def confirm(self,image_data):
+        config = {"configurable": {"thread_id": "1"}}
+        if image_data:
+            data_list=self.agent.get_state(config).values['image_data_list']
+            img_number=self.agent.get_state(config).values['image_number']
+            image_name=self.agent.get_state(config).values['image_name']
+            if not data_list:
+                    data_list=[]
+            data_list.append({'receipt_name':f'{img_number}_new_receipt.jpg',
+                'receipt_data':image_data})
+            self.agent.update_state(config,values={'image_data_list':data_list})
+            
+            
+            return data_list,image_name
+        
+  

requirements.txt

@@ -0,0 +1,10 @@
+ipython==9.0.1
+langchain==0.3.20
+langchain_core==0.3.41
+langchain_google_genai==2.0.11
+langgraph==0.3.5
+Pillow==11.1.0
+pydantic==2.10.6
+python-dotenv==1.0.1
+typing_extensions==4.12.2
+streamlit==1.43.0

streamlit.py

@@ -0,0 +1,326 @@
+import streamlit as st
+import json
+import copy
+from receipt_gen_agent import receipt_agent
+from PIL import Image
+
+from os import listdir
+from os.path import isfile, join
+agent=receipt_agent()
+def capture_or_upload_image():
+    """
+    Provides options to capture or upload an image using Streamlit.
+    
+    Returns:
+    PIL.Image.Image or None: The captured/uploaded image
+    """
+    st.sidebar.header("📸 Image Capture / Upload")
+    
+    # Choose method
+    image_source = st.sidebar.radio(
+        "Choose Image Source", 
+        ["Upload", "Camera Capture"],
+        key="image_source"
+    )
+    
+    # Image upload
+    if image_source == "Upload":
+        uploaded_file = st.sidebar.file_uploader(
+            "Upload Receipt Image", 
+            type=['png', 'jpg', 'jpeg', 'gif'],
+            help="Upload an image of the receipt",
+            key="uploaded_file"
+        )
+        
+        if uploaded_file is not None:
+            # Convert to PIL Image
+            pil_image = Image.open(uploaded_file)
+            # Store in session state
+            st.session_state.receipt_image = pil_image
+            return pil_image
+    
+    # Camera capture
+    else:
+        captured_image = st.camera_input(
+            "Take a picture of the receipt",
+            key="camera_input"
+        )
+        
+        if captured_image is not None:
+            # Convert to PIL Image
+            pil_image = Image.open(captured_image)
+            # Store in session state
+            st.session_state.receipt_image = pil_image
+            return pil_image
+    
+    # Return image from session state if available
+    return st.session_state.get("receipt_image", None)
+
+def parse_receipt_data(receipt_str):
+    """Parse the receipt dictionary from a string or dict."""
+    if isinstance(receipt_str, str):
+        try:
+            return json.loads(receipt_str.replace("'", '"'))
+        except json.JSONDecodeError:
+            return {}
+    return receipt_str
+
+def receipt_editor_app():
+    # Initialize session state variables if they don't exist
+    if "receipt_dict" not in st.session_state:
+        # Default initial receipt data
+        initial_receipt_str = {
+            "loc_name": "someplace",
+            "address": "somewhere",
+            "date": "someday",
+            "time": "sometime",
+            "items": [
+                {
+                    "name": "something",
+                    "price": 3.0,
+                    "quantity": 1
+                }
+            ],
+            "subtotal": 51.9,
+            "tax": 4.68,
+            "total": 56.58
+        }
+        st.session_state.receipt_dict = parse_receipt_data(initial_receipt_str)
+        st.session_state.edited_receipt = copy.deepcopy(st.session_state.receipt_dict)
+    
+    if "receipt_image" not in st.session_state:
+        st.session_state.receipt_image = None
+        
+    if "show_receipt_image" not in st.session_state:
+        st.session_state.show_receipt_image = True
+    
+    # Set page title and layout
+    st.set_page_config(page_title="Receipt Editor", layout="wide")
+
+    # Capture or upload image
+    receipt_image = capture_or_upload_image()
+    
+    if receipt_image is not None:
+        # Option to display image
+        show_image = st.sidebar.checkbox(
+            "Show Receipt Image", 
+            value=st.session_state.show_receipt_image,
+            key="show_image_checkbox"
+        )
+        st.session_state.show_receipt_image = show_image
+        
+        if show_image:
+            # Create two columns for image and details
+            col1, col2 = st.sidebar.columns([1, 1])
+            
+            with col1:
+                # Resize image to fit sidebar
+                resized_image = receipt_image.copy()
+                resized_image.thumbnail((600, 600))
+                st.image(resized_image, caption="Uploaded Receipt")
+            
+            with col2:
+                # Basic image details
+                st.write(f"Image Size: {receipt_image.size}")
+                st.write(f"Image Format: {receipt_image.format if hasattr(receipt_image, 'format') else 'Unknown'}")
+                st.write(f"Color Mode: {receipt_image.mode}")
+    
+    if st.sidebar.button("Process Image"):
+        initial_receipt_str = agent.receipt_gen(receipt_image)
+        # Parse initial receipt data
+        st.session_state.receipt_dict = parse_receipt_data(initial_receipt_str)
+    
+            
+        # Update edited receipt in session state
+        st.session_state.edited_receipt = copy.deepcopy(st.session_state.receipt_dict)
+        st.sidebar.success("Receipt processed!")
+        
+        
+    # Create a reference to session state for cleaner code
+    edited_receipt = st.session_state.edited_receipt
+    
+    # Title
+    st.title("Receipt Editor")
+    
+    # Basic Information Section
+    st.header("Basic Information")
+    col1, col2 = st.columns(2)
+    with col1:
+        edited_receipt['loc_name'] = st.text_input(
+            "Location Name", 
+            value=edited_receipt['loc_name'],
+            key="loc_name"
+        )
+        edited_receipt['date'] = st.text_input(
+            "Date", 
+            value=edited_receipt['date'],
+            key="date"
+        )
+    
+    with col2:
+        edited_receipt['address'] = st.text_area(
+            "Address", 
+            value=edited_receipt['address'],
+            key="address"
+        )
+        edited_receipt['time'] = st.text_input(
+            "Time", 
+            value=edited_receipt['time'],
+            key="time"
+        )
+    
+    # Items Section
+    st.header("Items")
+    
+    # Expandable section for items
+    with st.expander("Edit Items"):
+        # Allow adding new items
+        add_item = st.checkbox("Add New Item", key="add_item_checkbox")
+        if add_item:
+            with st.form(key="add_item_form"):
+                new_item_name = st.text_input("New Item Name", key="new_item_name")
+                new_item_price = st.number_input("New Item Price", min_value=0.0, step=0.01, key="new_item_price")
+                new_item_quantity = st.number_input("New Item Quantity", min_value=1, step=1, value=1, key="new_item_quantity")
+                
+                submit_button = st.form_submit_button(label="Add to Receipt")
+                if submit_button:
+                    new_item = {
+                        'name': new_item_name,
+                        'price': new_item_price,
+                        'quantity': new_item_quantity
+                    }
+                    edited_receipt['items'].append(new_item)
+                    st.session_state.edited_receipt = edited_receipt
+                    
+        
+    # Edit existing items
+    for i, item in enumerate(edited_receipt['items']):
+        with st.expander(f"Item {i+1}: {item['name']}"):
+            col1, col2, col3 = st.columns(3)
+            with col1:
+                item['name'] = st.text_input(
+                    f"Item {i+1} Name", 
+                    value=item['name'], 
+                    key=f"name_{i}"
+                )
+            with col2:
+                item['price'] = st.number_input(
+                    f"Item {i+1} Price", 
+                    min_value=0.0, 
+                    value=item['price'], 
+                    step=0.01, 
+                    key=f"price_{i}"
+                )
+            with col3:
+                item['quantity'] = st.number_input(
+                    f"Item {i+1} Quantity", 
+                    min_value=1, 
+                    value=item['quantity'], 
+                    step=1, 
+                    key=f"quantity_{i}"
+                )
+            
+            # Option to remove item
+            if st.button(f"Remove Item {i+1}", key=f"remove_{i}"):
+                del edited_receipt['items'][i]
+                st.session_state.edited_receipt = edited_receipt
+               
+
+    # Totals Section
+    st.header("Totals")
+    col1, col2, col3 = st.columns(3)
+    edited_receipt=st.session_state.edited_receipt
+    with col1:
+        edited_receipt['subtotal'] = st.number_input(
+            "Subtotal", 
+            min_value=0.0, 
+            value=edited_receipt['subtotal'], 
+            step=0.01,
+            key="subtotal"
+        )
+    
+    with col2:
+        edited_receipt['tax'] = st.number_input(
+            "Tax", 
+            min_value=0.0, 
+            value=edited_receipt['tax'], 
+            step=0.01,
+            key="tax"
+        )
+    
+    with col3:
+        edited_receipt['total'] = st.number_input(
+            "Total", 
+            min_value=0.0, 
+            value=edited_receipt['total'], 
+            step=0.01,
+            key="total"
+        )
+
+    # Recalculate totals button
+    if st.button("Recalculate Totals", key="recalc_button"):
+        edited_receipt=st.session_state.edited_receipt
+        # Automatically calculate subtotal
+        edited_receipt['subtotal'] = sum(item['price'] * item['quantity'] for item in edited_receipt['items'])
+        
+        # Recalculate total
+        edited_receipt['total'] = edited_receipt['subtotal'] + edited_receipt['tax']
+        
+        # Update session state
+        st.session_state.edited_receipt = edited_receipt
+    
+    # Display final edited receipt
+    st.header("Edited Receipt")
+    st.json(st.session_state.edited_receipt)
+
+    # Handle file saving
+
+    data_list = [f for f in listdir('new_receipt_data') if isfile(join('new_receipt_data', f))]
+    if not data_list:
+        st.session_state.data_list = []
+    else:
+        with open(f'new_receipt_data/{data_list[0]}', 'r') as openfile:
+            # Reading from json file
+            st.session_state.data_list = json.load(openfile)
+
+
+    st.session_state.image_name = f'{len(st.session_state.data_list)+1}_new_receipt.jpg'
+    
+    # Option to save or export
+    if st.button("Save Receipt Data", key="save_button"):
+        try:
+            # Create directories if they don't exist
+            
+            # Save receipt data
+            data_list=st.session_state.data_list
+            data_list.append({
+                'receipt_name':st.session_state.image_name,
+                'receipt_data': st.session_state.edited_receipt
+            })
+            
+            with open("new_receipt_data/receipt_data_list.json", "w") as f:
+                json.dump(data_list, f, indent=2)
+            
+            # Save image if available
+            if st.session_state.receipt_image:
+                save_path = f'new_images/{st.session_state.image_name}'
+                st.session_state.receipt_image.save(save_path)
+                st.success(f"Receipt data and image saved successfully!")
+            else:
+                st.success("Receipt data saved (no image available)")
+                
+        except Exception as e:
+            st.error(f"Error saving receipt: {str(e)}")
+
+    # Reset button to clear session state
+    if st.sidebar.button("Reset All", key="reset_button"):
+        # Clear specific session state values
+        for key in list(st.session_state.keys()):
+            del st.session_state[key]
+        
+
+def main():
+    receipt_editor_app()
+
+if __name__ == "__main__":
+    main()