Spaces:
Running
Running
| import functools | |
| import typing | |
| import aiohttp | |
| from langchain.docstore.document import Document | |
| from langchain import SerpAPIWrapper | |
| from src.utils_langchain import _chunk_sources, add_parser, _add_meta | |
| from urllib.parse import urlparse | |
| class H2OSerpAPIWrapper(SerpAPIWrapper): | |
| def get_search_documents(self, query, | |
| query_action=True, | |
| chunk=True, chunk_size=512, | |
| db_type='chroma', | |
| headsize=50, | |
| top_k_docs=-1): | |
| docs = self.run(query, headsize) | |
| chunk_sources = functools.partial(_chunk_sources, chunk=chunk, chunk_size=chunk_size, db_type=db_type) | |
| docs = chunk_sources(docs) | |
| # choose chunk type | |
| if query_action: | |
| docs = [x for x in docs if x.metadata['chunk_id'] >= 0] | |
| else: | |
| docs = [x for x in docs if x.metadata['chunk_id'] == -1] | |
| # get score assuming search results scale with ranking | |
| delta = 0.05 | |
| [x.metadata.update(score=0.1 + delta * x.metadata['chunk_id'] if x.metadata['chunk_id'] >= 0 else -1) for x in | |
| docs] | |
| # ensure see all results up to cutoff or mixing with non-web docs | |
| if top_k_docs >= 1: | |
| top_k_docs = max(top_k_docs, len(docs)) | |
| return docs, top_k_docs | |
| async def arun(self, query: str, headsize: int, **kwargs: typing.Any) -> list: | |
| """Run query through SerpAPI and parse result async.""" | |
| return self._process_response(await self.aresults(query), query, headsize) | |
| def run(self, query: str, headsize: int, **kwargs: typing.Any) -> list: | |
| """Run query through SerpAPI and parse result.""" | |
| return self._process_response(self.results(query), query, headsize) | |
| def _process_response(res: dict, query: str, headsize: int) -> list: | |
| try: | |
| return H2OSerpAPIWrapper.__process_response(res, query, headsize) | |
| except Exception as e: | |
| print("SERP search failed: %s" % str(e)) | |
| return [] | |
| def __process_response(res: dict, query: str, headsize: int) -> list: | |
| docs = [] | |
| res1 = SerpAPIWrapper._process_response(res) | |
| if res1: | |
| if isinstance(res1, str) and not res1.startswith('['): # avoid snippets | |
| docs += [Document(page_content='Web search result %s: ' % len(docs) + res1, | |
| metadata=dict(source='Web Search %s for %s' % (len(docs), query), score=0.0))] | |
| elif isinstance(res1, list): | |
| for x in res1: | |
| date = '' | |
| content = '' | |
| if 'source' in x: | |
| source = x['source'] | |
| content += '%s says' % source | |
| else: | |
| content = 'Web search result %s: ' % len(docs) | |
| if 'date' in x: | |
| date = x['date'] | |
| content += ' %s' % date | |
| if 'title' in x: | |
| content += ': %s' % x['title'] | |
| if 'snippet' in x: | |
| content += ': %s' % x['snippet'] | |
| if 'link' in x: | |
| link = x['link'] | |
| domain = urlparse(link).netloc | |
| font_size = 2 | |
| source_name = domain | |
| http_content = """<font size="%s"><a href="%s" target="_blank" rel="noopener noreferrer">%s</a></font>""" % ( | |
| font_size, link, source_name) | |
| source = 'Web Search %s' % len(docs) + \ | |
| ' from Date: %s Domain: %s Link: %s' % (date, domain, http_content) | |
| if date: | |
| content += ' around %s' % date | |
| content += ' according to %s' % domain | |
| else: | |
| source = 'Web Search %s for %s' % (len(docs), query) | |
| docs += [Document(page_content=content, metadata=dict(source=source, score=0.0))] | |
| if "knowledge_graph" in res.keys(): | |
| knowledge_graph = res["knowledge_graph"] | |
| title = knowledge_graph["title"] if "title" in knowledge_graph else "" | |
| if "description" in knowledge_graph.keys(): | |
| docs += [Document(page_content='Web search result %s: ' % len(docs) + knowledge_graph["description"], | |
| metadata=dict(source='Web Search %s with knowledge_graph description for %s' % ( | |
| len(docs), query), score=0.0))] | |
| for key, value in knowledge_graph.items(): | |
| if ( | |
| type(key) == str | |
| and type(value) == str | |
| and key not in ["title", "description"] | |
| and not key.endswith("_stick") | |
| and not key.endswith("_link") | |
| and not value.startswith("http") | |
| ): | |
| docs += [Document(page_content='Web search result %s: ' % len(docs) + f"{title} {key}: {value}.", | |
| metadata=dict( | |
| source='Web Search %s with knowledge_graph for %s' % (len(docs), query), | |
| score=0.0))] | |
| if "organic_results" in res.keys(): | |
| for org_res in res["organic_results"]: | |
| keys_to_try = ['snippet', 'snippet_highlighted_words', 'rich_snippet', 'rich_snippet_table', 'link'] | |
| for key in keys_to_try: | |
| if key in org_res.keys(): | |
| date = '' | |
| domain = '' | |
| link = '' | |
| snippet1 = '' | |
| if key != 'link': | |
| snippet1 = org_res[key] | |
| if 'date' in org_res.keys(): | |
| date = org_res['date'] | |
| snippet1 += ' on %s' % date | |
| else: | |
| date = 'unknown date' | |
| if 'link' in org_res.keys(): | |
| link = org_res['link'] | |
| domain = urlparse(link).netloc | |
| if key == 'link': | |
| # worst case, only url might have REST info | |
| snippet1 += ' Link at %s: <a href="%s">%s</a>' % (domain, link, domain) | |
| else: | |
| snippet1 += ' according to %s' % domain | |
| if snippet1: | |
| font_size = 2 | |
| source_name = domain | |
| http_content = """<font size="%s"><a href="%s" target="_blank" rel="noopener noreferrer">%s</a></font>""" % ( | |
| font_size, link, source_name) | |
| source = 'Web Search %s' % len(docs) + \ | |
| ' from Date: %s Domain: %s Link: %s' % (date, domain, http_content) | |
| domain_simple = domain.replace('www.', '').replace('.com', '') | |
| snippet1 = '%s says on %s: %s' % (domain_simple, date, snippet1) | |
| docs += [Document(page_content=snippet1, metadata=dict(source=source), score=0.0)] | |
| break | |
| if "buying_guide" in res.keys(): | |
| docs += [Document(page_content='Web search result %s: ' % len(docs) + res["buying_guide"], | |
| metadata=dict(source='Web Search %s with buying_guide for %s' % (len(docs), query)), | |
| score=0.0)] | |
| if "local_results" in res.keys() and "places" in res["local_results"].keys(): | |
| docs += [Document(page_content='Web search result %s: ' % len(docs) + res["local_results"]["places"], | |
| metadata=dict( | |
| source='Web Search %s with local_results_places for %s' % (len(docs), query)), | |
| score=0.0)] | |
| # add meta | |
| add_meta = functools.partial(_add_meta, headsize=headsize, parser='SERPAPI') | |
| add_meta(docs, query) | |
| return docs | |