| 797 | # --------------------------------------------------------------------------- |
| 798 | |
| 799 | class VertexAISearchChecker: |
| 800 | SEARCHLITE_URL = ( |
| 801 | "https://discoveryengine.googleapis.com/v1/projects/{project_id}" |
| 802 | "/locations/global/collections/default_collection" |
| 803 | "/engines/{app_id}/servingConfigs/default_search:searchLite" |
| 804 | ) |
| 805 | |
| 806 | def __init__(self, project_id: str, app_id: str, api_key: str, state: dict, |
| 807 | seed_mode: bool = False): |
| 808 | self.project_id = project_id |
| 809 | self.app_id = app_id |
| 810 | self.api_key = api_key |
| 811 | self.state = state |
| 812 | self.seed_mode = seed_mode |
| 813 | self.available = bool(project_id and app_id and api_key) |
| 814 | |
| 815 | def search(self, query: str, num: int = 10) -> list: |
| 816 | if not self.available: |
| 817 | return [] |
| 818 | url = self.SEARCHLITE_URL.format( |
| 819 | project_id=self.project_id, app_id=self.app_id, |
| 820 | ) |
| 821 | try: |
| 822 | resp = requests.post( |
| 823 | url, |
| 824 | params={"key": self.api_key}, |
| 825 | json={ |
| 826 | "query": query, |
| 827 | "pageSize": min(num, 25), |
| 828 | }, |
| 829 | headers={"Content-Type": "application/json"}, |
| 830 | timeout=20, |
| 831 | ) |
| 832 | if resp.status_code != 200: |
| 833 | body = resp.text[:300] if resp.text else "(empty)" |
| 834 | print(f" ERROR Vertex AI Search HTTP {resp.status_code}: {body}") |
| 835 | return [] |
| 836 | data = resp.json() |
| 837 | results = [] |
| 838 | for r in data.get("results", []): |
| 839 | doc = r.get("document", {}) |
| 840 | derived = doc.get("derivedStructData", {}) |
| 841 | link = derived.get("link", "") |
| 842 | title = derived.get("title", "") |
| 843 | snippet = "" |
| 844 | for s in derived.get("snippets", []): |
| 845 | if s.get("snippet"): |
| 846 | snippet = re.sub(r"<[^>]+>", "", s["snippet"]) |
| 847 | break |
| 848 | if title or link: |
| 849 | results.append({ |
| 850 | "title": title, |
| 851 | "link": link, |
| 852 | "snippet": snippet[:200], |
| 853 | }) |
| 854 | return results |
| 855 | except Exception as e: |
| 856 | print(f" ERROR Vertex AI Search: {e}") |