import gradio as gr import bs4 from langchain import hub from langchain_unstructured import UnstructuredLoader from langchain_core.documents import Document from typing_extensions import List, TypedDict from langchain_core.vectorstores import InMemoryVectorStore from langgraph.graph import START, StateGraph, MessagesState from langchain.chat_models import init_chat_model from langgraph.graph import END from langgraph.prebuilt import ToolNode, tools_condition from langchain_core.tools import tool from langchain_core.messages import SystemMessage, HumanMessage import getpass import os from langchain_huggingface import HuggingFaceEmbeddings import base64 import json import re import pytesseract import cv2 import numpy as np from pdf2image import convert_from_path # ---------- SETUP ---------- if not os.environ.get("GOOGLE_API_KEY"): os.environ["GOOGLE_API_KEY"] = getpass.getpass("Enter API key for Google Gemini: ") llm = init_chat_model("gemini-1.5-flash", model_provider="google_genai") embeddings = HuggingFaceEmbeddings( model_name="intfloat/multilingual-e5-large-instruct", model_kwargs={"device": 'cpu', "trust_remote_code": True} ) vector_store = InMemoryVectorStore(embeddings) prompt = hub.pull("rlm/rag-prompt") # ---------- RETRIEVAL TOOL ---------- @tool(response_format="content_and_artifact") def retrieve(query: str): """Retrieve information related to a query.""" retrieved_docs = vector_store.similarity_search(query, k=3) serialized = "\n\n".join( (f"Source: {doc.metadata}\nContent: {doc.page_content}") for doc in retrieved_docs ) return serialized, retrieved_docs # ---------- GRAPH FUNCTIONS FOR RAG ---------- def query_or_respond(state: MessagesState): """Generate tool call for retrieval or respond.""" llm_with_tools = llm.bind_tools([retrieve]) response = llm_with_tools.invoke(state["messages"]) return {"messages": [response]} tools = ToolNode([retrieve]) def generate(state: MessagesState): """Generate answer.""" recent_tool_messages = [] for message in reversed(state["messages"]): if message.type == "tool": recent_tool_messages.append(message) else: break tool_messages = recent_tool_messages[::-1] docs_content = "\n\n".join(doc.content for doc in tool_messages) print(f"retrieved docs: ", docs_content) system_message_content = ( "You are an assistant for question-answering tasks. " "Use the following pieces of retrieved context to answer " "the question. If you don't know the answer, say that you don't know. " "Use three sentences maximum and keep the answer concise." "\n\n" f"{docs_content}" ) conversation_messages = [ message for message in state["messages"] if message.type in ("human", "system") or (message.type == "ai" and not message.tool_calls) ] prompt = [SystemMessage(system_message_content)] + conversation_messages response = llm.invoke(prompt) return {"messages": [response]} # ---------- BUILD RAG GRAPH ---------- graph_builder = StateGraph(MessagesState) graph_builder.add_node(query_or_respond) graph_builder.add_node(tools) graph_builder.add_node(generate) graph_builder.set_entry_point("query_or_respond") graph_builder.add_conditional_edges( "query_or_respond", tools_condition, {END: END, "tools": "tools"}, ) graph_builder.add_edge("tools", "generate") graph_builder.add_edge("generate", END) rag_graph = graph_builder.compile() # ---------- FORM FILLING AGENTS ---------- class FormState(TypedDict): template_path: str source_path: str schema: dict filled_data: dict filled_image: str def find_bbox(file_path, prim_schema, alignment="down"): keys = list(prim_schema.keys()) img = cv2.imread(file_path) img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) data = pytesseract.image_to_data(img_rgb, output_type=pytesseract.Output.DICT) words = [] for i in range(len(data['text'])): if int(data['conf'][i]) > -1 and data['text'][i].strip(): words.append({ 'text': data['text'][i], 'left': data['left'][i], 'top': data['top'][i], 'width': data['width'][i], 'height': data['height'][i], 'conf': int(data['conf'][i]) }) words.sort(key=lambda w: (w['top'], w['left'])) lines = [] current_line = [] current_top = None if not words else words[0]['top'] for word in words: if current_line and abs(word['top'] - current_top) > 20: lines.append(current_line) current_line = [] current_line.append(word) current_top = word['top'] if current_line: lines.append(current_line) boxes = [] for line in lines: if line: full_text = ' '.join(w['text'] for w in line) left = min(w['left'] for w in line) top = min(w['top'] for w in line) right = max(w['left'] + w['width'] for w in line) bottom = max(w['top'] + w['height'] for w in line) boxes.append({ 'text': full_text, 'clean': full_text.lower().replace(" ", "").replace(":", ""), 'left': left, 'top': top, 'width': right - left, 'height': bottom - top }) boxes.sort(key=lambda b: (b['top'], b['left'])) schema = {} height, width = img.shape[:2] threshold = 40 for idx, box in enumerate(boxes): if box['clean'] in keys and box['clean'] not in schema: key = box['clean'] label_bbox_norm = [box['left'] / width, box['top'] / height, (box['left'] + box['width']) / width, (box['top'] + box['height']) / height] label_bbox_pixel = [box['left'], box['top'], box['left'] + box['width'], box['top'] + box['height']] input_bbox = None if alignment == "right": # Look for the next box to the right within the same line with tighter vertical alignment for next_idx in range(idx + 1, len(boxes)): next_box = boxes[next_idx] if (next_box['top'] >= box['top'] and next_box['top'] + next_box['height'] <= box['top'] + box['height'] + 10 and next_box['left'] > box['left'] + box['width'] and next_box['left'] < box['left'] + box['width'] + 300): input_bbox_norm = [next_box['left'] / width, next_box['top'] / height, (next_box['left'] + next_box['width']) / width, (next_box['top'] + next_box['height']) / height] input_bbox_pixel = [next_box['left'], next_box['top'], next_box['left'] + next_box['width'], next_box['top'] + next_box['height']] break else: # Default to "down" # Look for the next box below the key for next_idx in range(idx + 1, len(boxes)): next_box = boxes[next_idx] if next_box['top'] > box['top'] + box['height'] and abs(next_box['left'] - box['left']) < 50 and next_box['top'] - box['top'] < 100: input_bbox_norm = [next_box['left'] / width, next_box['top'] / height, (next_box['left'] + next_box['width']) / width, (next_box['top'] + next_box['height']) / height] input_bbox_pixel = [next_box['left'], next_box['top'], next_box['left'] + next_box['width'], next_box['top'] + next_box['height']] break if input_bbox is None: if alignment == "right": input_x = box['left'] + box['width'] + threshold input_y = box['top'] input_w = 200 # Adjusted to match typical input width in your image input_h = box['height'] else: # "down" input_x = box['left'] input_y = box['top'] + box['height'] + threshold input_w = box['width'] input_h = 20 input_bbox_norm = [input_x / width, input_y / height, (input_x + input_w) / width, (input_y + input_h) / height] input_bbox_pixel = [input_x, input_y, input_x + input_w, input_y + input_h] schema[key] = input_bbox_pixel return schema def convert_template_file(file_path): ext = os.path.splitext(file_path)[1].lower() if ext == ".pdf": images = convert_from_path(file_path, dpi=300) out_path = "template_converted.png" images[0].save(out_path, "PNG") return out_path else: raise ValueError("Unsupported template file format") def analyze_template(file_path: str) -> dict: if file_path.endswith("pdf"): file_path = convert_template_file(file_path) with open(file_path, "rb") as image_file: image_data = base64.b64encode(image_file.read()).decode("utf-8") message = { "role": "user", "content": [ { "type": "text", "text": "Analyse the following form and return just the JSON containing keys and values present in the image" "If te corresponding value is not present keep it as None." "Keep in mind that the keys cannot contain any spaces and should be lowercase. The output should be json loadable. ", }, { "type": "image", "source_type": "base64", "data": image_data, "mime_type": "image/jpeg", }, ] } response = llm.invoke([message]) out = response.text() match = re.search(r"\{.*\}", out, re.DOTALL) if match: json_string = match.group(0) # clean JSON schema = json.loads(json_string) else: schema = {} updated_schema = find_bbox(file_path, schema) position_schema = updated_schema return position_schema def extract_values(file_path: str, schema: dict) -> dict: filled_schema = {} schema_keys = list(schema.keys()) if file_path.endswith("txt"): with open(file_path, "r") as f: text = f.read() instr = f"Extract values from the source text for the following fields: {schema_keys}.Return a JSON with keys and extracted values.\n source text: {text}" message = { "role": "user", "content": [ { "type": "text", "text": instr, } ] } else: if file_path.endswith("pdf"): images = convert_from_path(file_path, dpi=300) out_path = "source_converted.png" images[0].save(out_path, "PNG") with open(out_path, "rb") as image_file: image_data = base64.b64encode(image_file.read()).decode("utf-8") if file_path.endswith("png") or file_path.endswith("jpg"): with open(file_path, "rb") as image_file: image_data = base64.b64encode(image_file.read()).decode("utf-8") text = f"Extract values from the image for the following fields: {schema_keys}.Return a JSON with keys and extracted values." message = { "role": "user", "content": [ { "type": "text", "text": text, }, { "type": "image", "source_type": "base64", "data": image_data, "mime_type": "image/jpeg", }, ] } response = llm.invoke([message]) out = response.text() match = re.search(r"\{.*\}", out, re.DOTALL) if match: json_string = match.group(0) # clean JSON filled = json.loads(json_string) else: filled = {} for key in schema: if key in filled: filled_schema[key] = filled[key] return filled_schema def fill_template(state: FormState): template_path = state["template_path"] position = state["schema"] filled_data = state["filled_data"] img = cv2.imread(template_path) for key, bbox in position.items(): if key in filled_data: x1, y1, x2, y2 = bbox text = filled_data[key] # Position text inside the box (slightly padded) cv2.putText(img, text, (x1+5, y2-5), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2) filled_path = template_path.replace(".jpg", "_filled.jpg").replace(".png", "_filled.png") cv2.imwrite(filled_path, img) return filled_path def analyze_node(state: FormState): schema = analyze_template(state["template_path"]) return {"schema": schema} def extract_node(state: FormState): filled = extract_values(state["source_path"], state["schema"]) return {"filled_data": filled} def fill_node(state: FormState): filled_image = fill_template(state) return {"filled_image": filled_image} # ---------- BUILD FORM FILLING GRAPH ---------- form_graph_builder = StateGraph(FormState) form_graph_builder.add_node("analyze", analyze_node) form_graph_builder.add_node("extract", extract_node) form_graph_builder.add_node("fill", fill_node) form_graph_builder.add_edge(START, "analyze") form_graph_builder.add_edge("analyze", "extract") form_graph_builder.add_edge("extract", "fill") form_graph_builder.add_edge("fill", END) form_graph = form_graph_builder.compile() # ---------- GRADIO APP FUNCTIONS ---------- def process_doc(file): loader = UnstructuredLoader( file_path=file.name, extract_images_in_pdf=True, languages=['ml', 'en'] ) docs = loader.load() vector_store.add_documents(documents=docs) return "✅ Document processed successfully! You can now ask questions." def ask_question(query, history): state = {"messages": [HumanMessage(content=query)]} response_text = "" for step in rag_graph.stream(state, stream_mode="values"): response_text = step["messages"][-1].content history.append((query, response_text)) return history, "" def process_form_filling(source, template): if not source or not template: return {"error": "Please upload both source and template files."}, None state = { "template_path": template.name, "source_path": source.name, "schema": {}, "filled_data": {}, "filled_image": "" } result = form_graph.invoke(state) return result["filled_data"], result["filled_image"] theme = gr.themes.Soft( primary_hue="blue", secondary_hue="gray", neutral_hue="slate", font=[gr.themes.GoogleFont("Inter"), gr.themes.GoogleFont("Source Sans Pro"), "ui-sans-serif", "system-ui"], font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "ui-monospace", "Consolas", "monospace"], ).set( body_background_fill="*neutral_50", body_background_fill_dark="*neutral_900", block_background_fill="*neutral_100", block_background_fill_dark="*neutral_800", block_border_width="1px", block_radius="8px", block_shadow="0 1px 3px rgba(0,0,0,0.1)", block_shadow_dark="0 1px 3px rgba(255,255,255,0.1)", button_primary_background_fill="*primary_500", button_primary_background_fill_hover="*primary_600", button_primary_text_color="white", button_secondary_background_fill="*neutral_200", button_secondary_background_fill_hover="*neutral_300", button_secondary_text_color="*neutral_800", input_background_fill="*neutral_50", input_background_fill_dark="*neutral_800", input_border_color="*neutral_200", input_border_color_dark="*neutral_700", panel_background_fill="*neutral_50", panel_background_fill_dark="*neutral_900", slider_color="*primary_500", ) with gr.Blocks(theme=theme, css=".gradio-container {max-width: 1200px !important; margin: auto;}") as demo: gr.Markdown( """ # 📑 Multi-lingual Doc RAG and Form Filling System """, elem_classes="text-center" ) with gr.Tabs(): with gr.Tab("📄 Document RAG"): with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Upload Document") upload_btn = gr.File(label="Select Document", file_types=[".pdf", ".txt", ".docx"], interactive=True) process_status = gr.Textbox(label="Processing Status", placeholder="Upload a document to start...", interactive=False) upload_btn.upload(process_doc, upload_btn, process_status) with gr.Column(scale=2): gr.Markdown("### Chat with Document") chatbot = gr.Chatbot(height=400, placeholder="Ask questions about your document here...") query = gr.Textbox(label="Your Question", placeholder="Type your question and press Enter...") query.submit(ask_question, [query, chatbot], [chatbot, query]) with gr.Tab("🖊️ Form Filling"): with gr.Row(): with gr.Column(): gr.Markdown("### Upload Files") source_upload = gr.File(label="Source File (Information Source)", file_types=[".jpg", ".png", ".txt", ".pdf"], interactive=True) template_upload = gr.File(label="Template File (Form to Fill)", file_types=[".jpg", ".png"], interactive=True) fill_btn = gr.Button("Process and Fill Form", variant="primary") with gr.Row(): with gr.Column(): gr.Markdown("### Extracted Data") output_json = gr.JSON(label="Filled Form Data (JSON)") with gr.Column(): gr.Markdown("### Filled Form Preview") output_image = gr.Image(label="Filled Form Image", interactive=False) fill_btn.click(process_form_filling, [source_upload, template_upload], [output_json, output_image]) gr.Markdown( """ --- *Note: For form filling the system currently expects all required fields to be completed under their corresponding keys.* """, elem_classes="text-center" ) demo.launch(server_name="0.0.0.0", server_port=7860)