Spaces:
Sleeping
Sleeping
| from .models import ParsedData, LlmExtraction | |
| from .database import SessionLocal | |
| import json | |
| class DataStore: | |
| def save_parsed_data(self, content: str) -> ParsedData: | |
| with SessionLocal() as session: | |
| parsed = ParsedData(Content=content) | |
| session.add(parsed) | |
| session.commit() | |
| session.refresh(parsed) | |
| return parsed | |
| def save_llm_extraction( | |
| self, parsed_data_id: int, json_data: dict, source_url: str | |
| ) -> LlmExtraction: | |
| with SessionLocal() as session: | |
| extraction = LlmExtraction( | |
| parsed_data_id=parsed_data_id, | |
| json_data=json.dumps(json_data), | |
| source_url=source_url, | |
| ) | |
| session.add(extraction) | |
| session.commit() | |
| session.refresh(extraction) | |
| return extraction | |
| def get_parsed_data(self, id: int) -> ParsedData | None: | |
| with SessionLocal() as session: | |
| return session.get(ParsedData, id) | |
| def get_all_parsed_data(self) -> list[ParsedData]: | |
| with SessionLocal() as session: | |
| return session.query(ParsedData).all() | |
| def get_llm_extractions_for_parsed( | |
| self, parsed_data_id: int | |
| ) -> list[LlmExtraction]: | |
| with SessionLocal() as session: | |
| return ( | |
| session.query(LlmExtraction) | |
| .filter_by(parsed_data_id=parsed_data_id) | |
| .all() | |
| ) | |