import functools import typing import aiohttp from langchain.docstore.document import Document from langchain.utilities 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) @staticmethod 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 [] @staticmethod 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 = """%s""" % ( 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: %s' % (domain, link, domain) else: snippet1 += ' according to %s' % domain if snippet1: font_size = 2 source_name = domain http_content = """%s""" % ( 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 def results(self, query: str) -> dict: # Fix non-thread-safe langchain swapping out sys directly. """Run query through SerpAPI and return the raw result.""" params = self.get_params(query) search = self.search_engine(params) res = search.get_dict() return res