Add cost tracking in Gradio interface

app.py

@@ -14,12 +14,14 @@ thinking_ladder_agent = ThinkingLadderAgent()
 synthesis_agent = SynthesisAgent()
 
 
-def main(query: str) -> str:
-    main_title, themes, _ = outline_agent.run(query)
-    sources, _ = research_agent.run(themes)
-    classified_sources, _ = thinking_ladder_agent.run(themes, sources)
-    synthesis, _ = synthesis_agent.run(main_title, themes, classified_sources)
-    return synthesis
+def main(query: str) -> tuple[str, float]:
+    main_title, themes, metrics_a = outline_agent.run(query)
+    sources, metrics_b = research_agent.run(themes)
+    classified_sources, metrics_c = thinking_ladder_agent.run(themes, sources)
+    synthesis, metrics_d = synthesis_agent.run(main_title, themes, classified_sources)
 
+    cost = sum(list(map(lambda metrics: metrics["cost"], [metrics_a, metrics_b, metrics_c, metrics_d])))
+    return synthesis, cost
 
-gr.Interface(fn=main, inputs=gr.Textbox(label="Query:"), outputs=gr.Textbox(label="Generated Synthesis:")).launch()
+
+gr.Interface(fn=main, inputs=gr.Textbox(label="Query:"), outputs=[gr.Textbox(label="Generated Synthesis:"), gr.Number(label="Cost ($):")]).launch()