Spaces:
Paused
Paused
| import os,re,json,gc,asyncio,calendar,traceback,faiss,torch,numpy as np,pdfplumber | |
| from datetime import date,datetime,timedelta | |
| from dateutil.relativedelta import relativedelta | |
| from sentence_transformers import SentenceTransformer,CrossEncoder | |
| from transformers import AutoTokenizer,AutoModelForCausalLM | |
| from fastapi import FastAPI,HTTPException,UploadFile,File,Form | |
| from fastapi.responses import FileResponse,StreamingResponse | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel | |
| from pypdf import PdfReader | |
| from reportlab.lib import colors | |
| from reportlab.lib.pagesizes import A4 | |
| from reportlab.lib.units import mm | |
| from reportlab.lib.enums import TA_LEFT,TA_JUSTIFY | |
| from reportlab.lib.styles import getSampleStyleSheet,ParagraphStyle | |
| from reportlab.platypus import SimpleDocTemplate,Paragraph,Spacer,Table,TableStyle,Image | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| BASE_PATH=os.environ.get("BASE_PATH","/app") | |
| MODEL_NAME=os.environ.get("MODEL_NAME","Qwen/Qwen2.5-3B-Instruct") | |
| SCHEMA_PATH=os.path.join(BASE_PATH,"json_files","schema.json") | |
| METRIC_PATH=os.path.join(BASE_PATH,"json_files","metric_catalog.json") | |
| COMPANY_DATA_PATH=os.path.join(BASE_PATH,"json_files","company_data.json") | |
| PROMPTS_PATH=os.path.join(BASE_PATH,"json_files","prompts.json") | |
| KNOWLEDGE_META_PATH=os.path.join(BASE_PATH,"report_knowledge_base","report_metadata.json") | |
| KNOWLEDGE_INDEX_PATH=os.path.join(BASE_PATH,"report_knowledge_base","report_vectors.index") | |
| OUTPUT_DIR=os.path.join(BASE_PATH,"output") | |
| os.makedirs(OUTPUT_DIR,exist_ok=True) | |
| app=FastAPI(title="Financial Intelligence API",version="1.0.0") | |
| app.add_middleware(CORSMiddleware,allow_origins=["*"],allow_credentials=True,allow_methods=["*"],allow_headers=["*"]) | |
| def load_json(path): | |
| with open(path,"r",encoding="utf-8") as f:return json.load(f) | |
| def save_json(data,path): | |
| directory=os.path.dirname(path) | |
| if directory:os.makedirs(directory,exist_ok=True) | |
| with open(path,"w",encoding="utf-8") as f:json.dump(data,f,indent=2,ensure_ascii=False) | |
| print("Loading company data...");company_data=load_json(COMPANY_DATA_PATH) | |
| print("Loading prompts...");PROMPTS=load_json(PROMPTS_PATH) | |
| print("Loading knowledge metadata...");knowledge_metadata=load_json(KNOWLEDGE_META_PATH) | |
| print("Loading embedding model...");embedding_model=SentenceTransformer("BAAI/bge-small-en-v1.5") | |
| print("Loading reranker...");reranker=CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2") | |
| print("Loading FAISS index...");knowledge_index=faiss.read_index(KNOWLEDGE_INDEX_PATH) | |
| print("Loading tokenizer...");tokenizer=AutoTokenizer.from_pretrained(MODEL_NAME,trust_remote_code=True) | |
| print("Loading Qwen model...") | |
| model=AutoModelForCausalLM.from_pretrained(MODEL_NAME,device_map="auto",torch_dtype=torch.float16,trust_remote_code=True) | |
| model.eval() | |
| print("Qwen model loaded successfully.") | |
| print("CUDA available:",torch.cuda.is_available()) | |
| if torch.cuda.is_available():print("GPU:",torch.cuda.get_device_name(0)) | |
| def safe_text(x): | |
| if isinstance(x,str):return x | |
| if isinstance(x,dict): | |
| for key in ["query","text","metric","name","field","value"]: | |
| if key in x and isinstance(x[key],(str,int,float)):return str(x[key]) | |
| return " ".join(str(value) for value in x.values()) | |
| if isinstance(x,list):return " ".join(safe_text(value) for value in x) | |
| return str(x or "") | |
| def is_date_value(value): | |
| if not isinstance(value,str):return False | |
| for fmt in ["%Y-%m-%d","%d-%m-%Y","%m/%d/%Y","%d/%m/%Y","%Y/%m/%d","%Y-%m-%d %H:%M:%S","%Y-%m-%dT%H:%M:%S"]: | |
| try:datetime.strptime(value,fmt);return True | |
| except ValueError:pass | |
| return False | |
| def detect_date_fields(rows): | |
| if not rows:return [] | |
| fields=[] | |
| for column in rows[0]: | |
| detected=any(is_date_value(row.get(column)) for row in rows if row.get(column)is not None) | |
| name_detected=any(word in column.lower() for word in ["date","time","year"]) | |
| if detected or name_detected:fields.append(column) | |
| return fields | |
| def extract_schema(data): | |
| output={} | |
| for table,rows in data.items(): | |
| if isinstance(rows,list) and rows and isinstance(rows[0],dict): | |
| sample=rows[0] | |
| output[table]={"columns":list(sample.keys()),"column_types":{col:type(val).__name__ for col,val in sample.items()},"numeric_fields":[col for col,val in sample.items() if isinstance(val,(int,float)) and not isinstance(val,bool)],"date_fields":detect_date_fields(rows)} | |
| elif isinstance(rows,dict): | |
| output[table]={"columns":list(rows.keys()),"column_types":{col:type(val).__name__ for col,val in rows.items()},"numeric_fields":[col for col,val in rows.items() if isinstance(val,(int,float)) and not isinstance(val,bool)],"date_fields":[col for col,val in rows.items() if is_date_value(val) or any(w in col.lower() for w in ["date","time","year"])]} | |
| return output | |
| schema=load_json(SCHEMA_PATH) if os.path.exists(SCHEMA_PATH) else (lambda s:save_json(s,SCHEMA_PATH) or s)(extract_schema(company_data)) | |
| def build_metric_catalog(schema): | |
| catalog={} | |
| for table,info in schema.items(): | |
| for field in info["columns"]: | |
| if field in info.get("date_fields",[]):continue | |
| catalog[f"{table}*{field}"]={"table":table,"field":field,"display_name":field.replace("_"," ").title(),"field_type":"number" if field in info.get("numeric_fields",[]) else info["column_types"].get(field)} | |
| return catalog | |
| metric_catalog=load_json(METRIC_PATH) if os.path.exists(METRIC_PATH) else (lambda c:save_json(c,METRIC_PATH) or c)(build_metric_catalog(schema)) | |
| metric_names=list(metric_catalog.keys()) | |
| metric_texts=[f"{metric_catalog[name]['table']} {metric_catalog[name]['display_name']}".lower() for name in metric_names] | |
| print("Building metric embeddings...") | |
| metric_embeddings=embedding_model.encode(metric_texts,normalize_embeddings=True).astype("float32") | |
| metric_index=faiss.IndexFlatIP(metric_embeddings.shape[1]) | |
| metric_index.add(metric_embeddings) | |
| def cleanup_gpu(): | |
| gc.collect() | |
| if torch.cuda.is_available():torch.cuda.empty_cache() | |
| llm_lock=asyncio.Lock() | |
| def _llm_json_sync(messages,max_tokens=1500): | |
| inputs=None;output=None | |
| try: | |
| prompt=tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True) | |
| inputs=tokenizer(prompt,return_tensors="pt",truncation=True,max_length=tokenizer.model_max_length) | |
| device=next(model.parameters()).device | |
| inputs={k:v.to(device) for k,v in inputs.items()} | |
| with torch.inference_mode(): | |
| output=model.generate(**inputs,max_new_tokens=max_tokens,do_sample=False,use_cache=True,pad_token_id=tokenizer.eos_token_id) | |
| generated=output[0][inputs["input_ids"].shape[-1]:] | |
| text=re.sub(r"```(?:json)?","",tokenizer.decode(generated,skip_special_tokens=True),flags=re.IGNORECASE).replace("```","").strip() | |
| try:return json.loads(text) | |
| except Exception:pass | |
| for pat in [r"\{.*\}",r"\[.*\]"]: | |
| if match:=re.search(pat,text,re.DOTALL): | |
| try:return json.loads(match.group(0)) | |
| except Exception:pass | |
| print("WARNING: Invalid LLM JSON:",text[:2000]) | |
| return {} | |
| except Exception: | |
| traceback.print_exc();return {} | |
| finally: | |
| del inputs,output | |
| cleanup_gpu() | |
| async def llm_json(messages,max_tokens=1500): | |
| async with llm_lock:return await asyncio.to_thread(_llm_json_sync,messages,max_tokens) | |
| def parse_financial_query(query): | |
| system="""Return JSON only: | |
| {"queries":[],"operations":[]} | |
| Extract independent financial data requests as plain strings and required operations. | |
| Do not extract database fields or answer the question. | |
| Operations: | |
| sum,count,average,compare,trend,percentage,group_by,status,explain_change | |
| Do not create queries for operations. | |
| Keep time information inside relevant queries.""" | |
| result=_llm_json_sync([{"role":"system","content":system},{"role":"user","content":query}],300) | |
| if not isinstance(result,dict):result={} | |
| queries=result.get("queries",[]) | |
| if not isinstance(queries,list):queries=[query] | |
| queries=[safe_text(item).strip() for item in queries if safe_text(item).strip()] or [query] | |
| operations=result.get("operations",[]) | |
| if not isinstance(operations,list):operations=[] | |
| return {"queries":queries,"operations":[safe_text(item).lower().strip() for item in operations]} | |
| def metric_embedding_rank(query,k=5): | |
| query=safe_text(query).strip() | |
| if not query:return [] | |
| embedding=embedding_model.encode([query],normalize_embeddings=True).astype("float32") | |
| k=min(k,len(metric_names)) | |
| if k<=0:return [] | |
| scores,indices=metric_index.search(embedding,k) | |
| results=[] | |
| for score,index in zip(scores[0],indices[0]): | |
| if index==-1:continue | |
| metric_name=metric_names[index];metric=metric_catalog[metric_name] | |
| results.append({"metric":metric_name,"table":metric["table"],"field":metric["field"],"display_name":metric["display_name"],"field_type":metric["field_type"],"embedding_score":float(score)}) | |
| return results | |
| def metric_semantic_rank(query,k=5): | |
| candidates=metric_embedding_rank(query,k) | |
| if not candidates:return [] | |
| pairs=[[query.lower(),f"{item['table']} {item['display_name']}".lower()] for item in candidates] | |
| scores=reranker.predict(pairs) | |
| for item,score in zip(candidates,scores):item["semantic_score"]=float(score) | |
| return sorted(candidates,key=lambda item:item["semantic_score"],reverse=True) | |
| def retrieve_metrics(parsed): | |
| return {safe_text(query):metric_semantic_rank(query,5) for query in parsed.get("queries",[])} | |
| MONTHS={name.lower():index for index,name in enumerate(calendar.month_name) if name} | |
| def resolve_date_range(query): | |
| query=safe_text(query).lower();today=date.today() | |
| if "this month" in query: | |
| start=today.replace(day=1);end=today.replace(day=calendar.monthrange(today.year,today.month)[1]) | |
| return {"start_date":start.isoformat(),"end_date":end.isoformat()} | |
| if "last month" in query: | |
| end=today.replace(day=1)-timedelta(days=1) | |
| return {"start_date":end.replace(day=1).isoformat(),"end_date":end.isoformat()} | |
| if match:=re.search(r"last\s+(\d+)\s+months?",query): | |
| count=max(1,int(match.group(1)));start=today.replace(day=1)-relativedelta(months=count-1);end=today.replace(day=calendar.monthrange(today.year,today.month)[1]) | |
| return {"start_date":start.isoformat(),"end_date":end.isoformat()} | |
| if match:=re.search(r"\b(january|february|march|april|may|june|july|august|september|october|november|december)\s+(\d{4})\s+to\s+(january|february|march|april|may|june|july|august|september|october|november|december)\s+(\d{4})\b",query): | |
| sm,sy,em,ey=MONTHS[match.group(1)],int(match.group(2)),MONTHS[match.group(3)],int(match.group(4)) | |
| return {"start_date":date(sy,sm,1).isoformat(),"end_date":date(ey,em,calendar.monthrange(ey,em)[1]).isoformat()} | |
| if match:=re.search(r"\b(\d{2})/(\d{4})\s+to\s+(\d{2})/(\d{4})\b",query): | |
| sm,sy,em,ey=map(int,match.groups()) | |
| if 1<=sm<=12 and 1<=em<=12:return {"start_date":date(sy,sm,1).isoformat(),"end_date":date(ey,em,calendar.monthrange(ey,em)[1]).isoformat()} | |
| if match:=re.search(r"from\s+(\d{2})/(\d{2})/(\d{4})\s+to\s+(\d{2})/(\d{2})/(\d{4})",query): | |
| d1,m1,y1,d2,m2,y2=map(int,match.groups()) | |
| return {"start_date":date(y1,m1,d1).isoformat(),"end_date":date(y2,m2,d2).isoformat()} | |
| if match:=re.search(r"from\s+(\d{4})\s+to\s+(\d{4})",query):return {"start_date":f"{match.group(1)}-01-01","end_date":f"{match.group(2)}-12-31"} | |
| if match:=re.search(r"\b(january|february|march|april|may|june|july|august|september|october|november|december)\s+(\d{4})\b",query): | |
| m,y=MONTHS[match.group(1)],int(match.group(2)) | |
| return {"start_date":date(y,m,1).isoformat(),"end_date":date(y,m,calendar.monthrange(y,m)[1]).isoformat()} | |
| if match:=re.search(r"\b(\d{2})/(\d{4})\b",query): | |
| m,y=map(int,match.groups()) | |
| if 1<=m<=12:return {"start_date":date(y,m,1).isoformat(),"end_date":date(y,m,calendar.monthrange(y,m)[1]).isoformat()} | |
| return None | |
| def build_intent(parsed,results): | |
| output=[] | |
| for query,metrics in results.items(): | |
| if not metrics:continue | |
| metric=metrics[0];date_range=resolve_date_range(query) | |
| output.append({"query":query,"metric":{"table":metric["table"],"field":metric["field"],"name":metric["display_name"],"metric_id":metric["metric"],"semantic_score":metric.get("semantic_score",0)},"operations":parsed.get("operations",[]),"filters":{"date_range":date_range} if date_range else {}}) | |
| return output | |
| def get_date_field(row): | |
| return next((column for column in row if any(word in column.lower() for word in ["date","paid_on","approved_on","created_at","timestamp","month"])),None) | |
| def retrieve_answer_data(intent,data): | |
| table=intent["metric"]["table"];field=intent["metric"]["field"];date_range=intent.get("filters",{}).get("date_range");rows=data.get(table,[]) | |
| if isinstance(rows,dict):rows=[rows] | |
| if not isinstance(rows,list):return [] | |
| result=[] | |
| for row in rows: | |
| if not isinstance(row,dict):continue | |
| date_field=get_date_field(row);row_date=str(row.get(date_field))[:10] if date_field and row.get(date_field)is not None else None | |
| if date_range and (not row_date or not(date_range["start_date"]<=row_date<=date_range["end_date"])):continue | |
| result.append({"date":row_date,intent["metric"]["name"]:row.get(field)}) | |
| return result | |
| def calculate_metric(intent,data_result): | |
| metric=intent["metric"]["name"];operations=intent.get("operations",[]) | |
| values=[{"date":row.get("date"),"value":row.get(metric)} for row in data_result if isinstance(row,dict) and isinstance(row.get(metric),(int,float)) and not isinstance(row.get(metric),bool)] | |
| if not values:return {"metric":metric,"status":"no numeric data","data":[]} | |
| result={"metric":metric,"data":values} | |
| if "sum" in operations:result={"metric":metric,"sum":sum(x["value"] for x in values),"data":values} | |
| if "count" in operations:result={"metric":metric,"count":len(values),"data":values} | |
| if "average" in operations:result={"metric":metric,"average":round(sum(x["value"] for x in values)/len(values),2),"data":values} | |
| if "percentage" in operations and len(values)>=2: | |
| previous,current=values[-2],values[-1];result["percentage_change"]=round(((current["value"]-previous["value"])/previous["value"]*100) if previous["value"] else 0,2) | |
| if "compare" in operations and len(values)>=2: | |
| previous,current=values[-2],values[-1];difference=current["value"]-previous["value"] | |
| result["comparison"]={"current_period":current,"previous_period":previous,"current_value":current["value"],"previous_value":previous["value"],"difference":difference,"percentage_change":round((difference/previous["value"]*100) if previous["value"] else 0,2)} | |
| if "trend" in operations and len(values)>=2:result={"metric":metric,"periods":values,"overall_change":values[-1]["value"]-values[0]["value"]} | |
| return result | |
| def retrieve_history(query,k=5): | |
| query=safe_text(query).strip() | |
| if not query:return [] | |
| embedding=embedding_model.encode([query],normalize_embeddings=True).astype("float32");k=min(k,len(knowledge_metadata)) | |
| if k<=0:return [] | |
| scores,indices=knowledge_index.search(embedding,k);candidates=[] | |
| for score,index in zip(scores[0],indices[0]): | |
| if index==-1:continue | |
| item=knowledge_metadata[index] if isinstance(knowledge_metadata[index],dict) else {};content=safe_text(item.get("content","")) | |
| if len(content.split())>=5:candidates.append({"id":int(index),"content":content,"embedding_score":float(score)}) | |
| if not candidates:return [] | |
| scores=reranker.predict([[query.lower(),x["content"].lower()] for x in candidates]) | |
| for item,score in zip(candidates,scores):item["semantic_score"]=float(score) | |
| candidates.sort(key=lambda x:x["semantic_score"],reverse=True) | |
| return [{"content":x["content"],"embedding_score":x["embedding_score"],"semantic_score":x["semantic_score"]} for x in candidates[:3]] | |
| def execute_calculations(intents): | |
| return [{"query":intent,"calculation":calculate_metric(intent,retrieve_answer_data(intent,company_data))} for intent in intents] | |
| def prepare_context(parsed,intents): | |
| context=[] | |
| for intent in intents: | |
| rows=retrieve_answer_data(intent,company_data) | |
| context.append({"query":intent,"intent":intent,"current_data":rows,"calculation":calculate_metric(intent,rows),"historical_knowledge":retrieve_history(intent["query"])}) | |
| return context | |
| def generate_financial_response(context): | |
| system=PROMPTS.get("financial_report",{}).get("system_prompt","\nReturn a structured financial report as JSON with:\n\nreport_title,\nreport_type,\nexecutive_summary,\nsections,\ntables,\ncharts,\nconclusion.\n") | |
| result=_llm_json_sync([{"role":"system","content":system},{"role":"user","content":json.dumps(context,default=str,ensure_ascii=False)}],1800) | |
| if not isinstance(result,dict):result={} | |
| defaults={"report_title":"Financial Analysis Report","report_type":"financial_report","executive_summary":"","sections":[],"tables":[],"charts":[],"conclusion":""} | |
| for key,value in defaults.items(): | |
| if key not in result:result[key]=value | |
| for key in ["sections","tables","charts"]: | |
| if not isinstance(result.get(key),list):result[key]=[] | |
| return result | |
| def extract_knowledge_from_pdf(path): | |
| with pdfplumber.open(path) as pdf:text="\n".join(p.extract_text(layout=True) or "" for p in pdf.pages) | |
| rows=[] | |
| for i in range(0,len(text),2500): | |
| res=_llm_json_sync([{"role":"system","content":PROMPTS["report_extraction"]["system"]},{"role":"user","content":text[i:i+2500]}],600) | |
| rows.extend(res if isinstance(res,list) else next((res[k] for k in ["rows","data","results","knowledge"] if isinstance(res,dict) and isinstance(res.get(k),list)),[])) | |
| clean=[];seen=set() | |
| for r in rows: | |
| if isinstance(r,dict): | |
| c=safe_text(r.get("content","")).strip() | |
| if c: | |
| key=(r.get("report_type"),safe_text(r.get("section","")).strip(),c) | |
| if key not in seen:seen.add(key);clean.append(r) | |
| existing={(x.get("report_type"),x.get("section"),x.get("content")) for x in knowledge_metadata} | |
| new_rows=[r for r in clean if (r.get("report_type"),r.get("section"),r.get("content")) not in existing] | |
| if new_rows: | |
| vectors=np.array(embedding_model.encode([x["content"] for x in new_rows]),dtype="float32");faiss.normalize_L2(vectors);knowledge_index.add(vectors);knowledge_metadata.extend(new_rows);faiss.write_index(knowledge_index,KNOWLEDGE_INDEX_PATH);save_json(knowledge_metadata,KNOWLEDGE_META_PATH) | |
| return {"extracted":len(rows),"cleaned":len(clean),"added":len(new_rows),"total_knowledge":knowledge_index.ntotal} | |
| def clean_text(text):return re.sub(r"[\*#`_]","",safe_text(text)).strip() | |
| def pdf_styles(): | |
| styles=getSampleStyleSheet() | |
| return {"title":ParagraphStyle("financial_title",parent=styles["Title"],fontName="Helvetica-Bold",fontSize=21,leading=25,textColor=colors.HexColor("#1A365D"),alignment=TA_LEFT,spaceAfter=6),"sub":ParagraphStyle("financial_subtitle",parent=styles["Normal"],fontSize=9,textColor=colors.HexColor("#64748B"),spaceAfter=12),"head":ParagraphStyle("financial_heading",parent=styles["Heading2"],fontName="Helvetica-Bold",fontSize=14,leading=18,textColor=colors.HexColor("#2E6F9E"),spaceBefore=12,spaceAfter=7),"body":ParagraphStyle("financial_body",parent=styles["BodyText"],fontSize=10,leading=15,textColor=colors.HexColor("#334155"),alignment=TA_JUSTIFY,spaceAfter=8)} | |
| def add_page_number(canvas,document): | |
| canvas.saveState();canvas.setFont("Helvetica",8);canvas.setFillColor(colors.HexColor("#64748B"));canvas.drawRightString(A4[0]-18*mm,10*mm,f"Page {document.page}");canvas.restoreState() | |
| def add_pdf_table(elements,table_data,styles): | |
| headers=table_data.get("headers",[]);rows=table_data.get("rows",[]) | |
| if not headers or not rows:return | |
| header_style=ParagraphStyle("financial_table_header",parent=styles["body"],fontName="Helvetica-Bold",textColor=colors.white,fontSize=9) | |
| data=[[Paragraph(clean_text(v),header_style) for v in headers]] | |
| for row in rows: | |
| if isinstance(row,list):data.append([Paragraph(clean_text(v),styles["body"]) for v in row]) | |
| if len(data)<2:return | |
| table=Table(data,repeatRows=1,hAlign="LEFT") | |
| table.setStyle(TableStyle([("BACKGROUND",(0,0),(-1,0),colors.HexColor("#1A365D")),("GRID",(0,0),(-1,-1),0.4,colors.HexColor("#CBD5E1")),("ROWBACKGROUNDS",(0,1),(-1,-1),[colors.white,colors.HexColor("#F8FAFC")]),("VALIGN",(0,0),(-1,-1),"TOP"),("LEFTPADDING",(0,0),(-1,-1),7),("RIGHTPADDING",(0,0),(-1,-1),7),("TOPPADDING",(0,0),(-1,-1),6),("BOTTOMPADDING",(0,0),(-1,-1),6)])) | |
| if table_data.get("title"):elements.append(Paragraph(clean_text(table_data["title"]),styles["head"])) | |
| elements.extend([table,Spacer(1,10)]) | |
| def create_chart(chart,index): | |
| data=chart.get("data",[]) if isinstance(chart,dict) else [];labels=[];values=[] | |
| for item in data: | |
| if isinstance(item,dict): | |
| try:labels.append(str(item.get("label","")));values.append(float(item.get("value",0))) | |
| except Exception:continue | |
| if len(values)<2:return None | |
| chart_type=str(chart.get("type","line")).lower();path=os.path.join(OUTPUT_DIR,f"chart_{index}.png") | |
| plt.figure(figsize=(8,4.5)) | |
| if chart_type=="bar":plt.bar(labels,values) | |
| elif chart_type=="pie":plt.pie(values,labels=labels,autopct="%1.1f%%") | |
| elif chart_type=="histogram":plt.bar(labels,values,width=1) | |
| else:plt.plot(labels,values,marker="o") | |
| plt.title(clean_text(chart.get("title","Financial Overview")));plt.xticks(rotation=30,ha="right");plt.tight_layout();plt.savefig(path,dpi=180,bbox_inches="tight");plt.close() | |
| return path | |
| def create_pdf(response,path): | |
| styles=pdf_styles();document=SimpleDocTemplate(path,pagesize=A4,leftMargin=18*mm,rightMargin=18*mm,topMargin=18*mm,bottomMargin=18*mm) | |
| elements=[Paragraph(clean_text(response.get("report_title","Financial Analysis Report")),styles["title"]),Paragraph(datetime.now().strftime("%B %d, %Y"),styles["sub"])] | |
| if summary:=response.get("executive_summary",""):elements.extend([Paragraph("Executive Summary",styles["head"]),Paragraph(clean_text(summary),styles["body"])]) | |
| for section in response.get("sections",[]): | |
| if isinstance(section,dict): | |
| if section.get("title"):elements.append(Paragraph(clean_text(section["title"]),styles["head"])) | |
| if section.get("content"):elements.append(Paragraph(clean_text(section["content"]),styles["body"])) | |
| for table_data in response.get("tables",[])[:2]: | |
| if isinstance(table_data,dict):add_pdf_table(elements,table_data,styles) | |
| chart_paths=[] | |
| for index,chart in enumerate(response.get("charts",[])[:2]): | |
| if isinstance(chart,dict) and (chart_path:=create_chart(chart,index)): | |
| chart_paths.append(chart_path);elements.extend([Paragraph(clean_text(chart.get("title","Financial Overview")),styles["head"]),Image(chart_path,width=165*mm,height=90*mm),Spacer(1,8)]) | |
| if conclusion:=response.get("conclusion",""):elements.extend([Paragraph("Conclusion",styles["head"]),Paragraph(clean_text(conclusion),styles["body"])]) | |
| document.build(elements,onFirstPage=add_page_number,onLaterPages=add_page_number) | |
| for p in chart_paths: | |
| try:os.remove(p) | |
| except Exception:pass | |
| return path | |
| def extract_pdf_text(path): | |
| reader=PdfReader(path) | |
| return "\n".join(page.extract_text() or "" for page in reader.pages).strip() | |
| class PromptRequest(BaseModel): | |
| prompt:str | |
| def progress_event(message,step="progress",data=None): | |
| payload={"step":step,"message":message} | |
| if isinstance(data,dict):payload.update(data) | |
| return f"data: {json.dumps(payload,default=str,ensure_ascii=False)}\n\n" | |
| async def health():return {"status":"ok","model":MODEL_NAME,"gpu":torch.cuda.is_available()} | |
| async def health_check(): | |
| gpu_name=torch.cuda.get_device_name(0) if torch.cuda.is_available() else None | |
| return {"status":"healthy","gpu":torch.cuda.is_available(),"gpu_name":gpu_name,"model":MODEL_NAME} | |
| async def generate_report(req:PromptRequest): | |
| try: | |
| query=re.sub(r",+"," ",req.prompt).strip() | |
| if not query:raise HTTPException(400,"prompt is required") | |
| parsed=parse_financial_query(query);metrics=retrieve_metrics(parsed);intents=build_intent(parsed,metrics) | |
| if not intents:raise HTTPException(400,"Unable to identify a financial metric") | |
| context=prepare_context(parsed,intents);calculations=execute_calculations(intents) | |
| if not any(item["current_data"] for item in context):raise HTTPException(404,"No financial data found matching request criteria") | |
| response=generate_financial_response({"user_query":query,"financial_context":context,"calculations":calculations}) | |
| if not response.get("executive_summary") and not response.get("sections"):raise HTTPException(500,"Unable to generate financial report") | |
| path=os.path.join(OUTPUT_DIR,"financial_report.pdf");create_pdf(response,path) | |
| return FileResponse(path,media_type="application/pdf",filename="financial_report.pdf") | |
| except HTTPException:raise | |
| except Exception as error: | |
| traceback.print_exc();raise HTTPException(500,f"{type(error).__name__}: {str(error)}") | |
| async def generate_report_stream(req:PromptRequest): | |
| async def event_stream(): | |
| try: | |
| query=re.sub(r",+"," ",req.prompt).strip() | |
| if not query:yield progress_event("Prompt is required","error");return | |
| yield progress_event("Understanding your request","parsing") | |
| parsed=await asyncio.to_thread(parse_financial_query,query) | |
| yield progress_event("Identifying the relevant financial information","metric") | |
| metrics=await asyncio.to_thread(retrieve_metrics,parsed);intents=await asyncio.to_thread(build_intent,parsed,metrics) | |
| if not intents:yield progress_event("I couldn't identify the requested financial information","error");return | |
| yield progress_event("Retrieving financial data","data_retrieval") | |
| context=await asyncio.to_thread(prepare_context,parsed,intents) | |
| if not any(item["current_data"] for item in context):yield progress_event("No financial data was found for your request","error");return | |
| yield progress_event("Calculating financial results","calculation") | |
| calculations=await asyncio.to_thread(execute_calculations,intents) | |
| yield progress_event("Analyzing the results","analysis") | |
| response=await asyncio.to_thread(generate_financial_response,{"user_query":query,"financial_context":context,"calculations":calculations}) | |
| if not response.get("executive_summary") and not response.get("sections"):yield progress_event("Unable to generate the financial report","error");return | |
| if summary:=safe_text(response.get("executive_summary","")).strip():yield progress_event(summary,"analysis_result") | |
| yield progress_event("Generating PDF report","pdf_generation") | |
| path=os.path.join(OUTPUT_DIR,"financial_report.pdf");await asyncio.to_thread(create_pdf,response,path) | |
| if not os.path.exists(path):yield progress_event("PDF generation failed","error");return | |
| data={"filename":"financial_report.pdf","download_endpoint":"/download-pdf"} | |
| yield progress_event("PDF generated successfully. Starting download.","pdf_ready",data) | |
| yield progress_event("Download ready","complete",data) | |
| except Exception: | |
| traceback.print_exc();yield progress_event("An error occurred while generating the report","error") | |
| return StreamingResponse(event_stream(),media_type="text/event-stream",headers={"Cache-Control":"no-cache","Connection":"keep-alive","X-Accel-Buffering":"no"}) | |
| async def generate_report_with_pdf(prompt:str=Form(...),file:UploadFile=File(...)): | |
| try: | |
| query=re.sub(r",+"," ",prompt).strip() | |
| if not query:raise HTTPException(400,"prompt is required") | |
| filename=os.path.basename(file.filename);upload_path=os.path.join(OUTPUT_DIR,f"uploaded_{filename}");contents=await file.read() | |
| with open(upload_path,"wb") as output_file:output_file.write(contents) | |
| pdf_text=extract_pdf_text(upload_path) | |
| if not pdf_text:raise HTTPException(400,"Unable to extract text from PDF") | |
| pdf_text=pdf_text[:12000];parsed=parse_financial_query(query);metrics=retrieve_metrics(parsed);intents=build_intent(parsed,metrics);context=prepare_context(parsed,intents) if intents else [] | |
| context.append({"query":query,"intent":{},"current_data":[],"calculation":{},"historical_knowledge":[{"content":pdf_text}]}) | |
| response=generate_financial_response({"user_query":query,"financial_context":context,"calculations":execute_calculations(intents) if intents else [],"uploaded_document_text":pdf_text}) | |
| if not response.get("executive_summary") and not response.get("sections"):raise HTTPException(500,"Unable to generate financial report") | |
| path=os.path.join(OUTPUT_DIR,"financial_report.pdf");create_pdf(response,path) | |
| return FileResponse(path,media_type="application/pdf",filename="financial_report.pdf") | |
| except HTTPException:raise | |
| except Exception as error: | |
| traceback.print_exc();raise HTTPException(500,f"{type(error).__name__}: {str(error)}") | |
| async def upload_pdf(file:UploadFile=File(...)): | |
| try: | |
| filename=os.path.basename(file.filename);path=os.path.join(OUTPUT_DIR,f"uploaded_{filename}");contents=await file.read() | |
| with open(path,"wb") as output_file:output_file.write(contents) | |
| history=await asyncio.to_thread(extract_knowledge_from_pdf,path) | |
| return {"status":"success","filename":file.filename,"path":path,"knowledge_extraction":history} | |
| except Exception as error: | |
| traceback.print_exc();raise HTTPException(500,f"{type(error).__name__}: {str(error)}") | |
| async def download_pdf(): | |
| path=os.path.join(OUTPUT_DIR,"financial_report.pdf") | |
| if not os.path.exists(path):raise HTTPException(404,"No report available") | |
| return FileResponse(path,media_type="application/pdf",filename="financial_report.pdf") | |
| async def shutdown_event(): | |
| gc.collect() | |
| if torch.cuda.is_available():torch.cuda.empty_cache() |