npcsh 0.3.31__py3-none-any.whl → 1.0.0__py3-none-any.whl

npcsh/search.py

@@ -1,248 +0,0 @@
-# search.py
-
-import requests
-import os
-
-from bs4 import BeautifulSoup
-from duckduckgo_search import DDGS
-from duckduckgo_search.exceptions import DuckDuckGoSearchException
-
-try:
-    from googlesearch import search
-except:
-    pass
-from typing import List, Dict, Any, Optional, Union
-import numpy as np
-import json
-
-try:
-    from sentence_transformers import util, SentenceTransformer
-except:
-    pass
-
-
-def search_perplexity(
-    query: str,
-    api_key: str = None,
-    model: str = "sonar",
-    max_tokens: int = 400,
-    temperature: float = 0.2,
-    top_p: float = 0.9,
-):
-    if api_key is None:
-        api_key = os.environ["PERPLEXITY_API_KEY"]
-    # print("api_key", api_key)
-    url = "https://api.perplexity.ai/chat/completions"
-    payload = {
-        "model": "sonar",
-        "messages": [
-            {"role": "system", "content": "Be precise and concise."},
-            {"role": "user", "content": query},
-        ],
-        "max_tokens": max_tokens,
-        "temperature": temperature,
-        "top_p": top_p,
-        "return_images": False,
-        "return_related_questions": False,
-        "search_recency_filter": "month",
-        "top_k": 0,
-        "stream": False,
-        "presence_penalty": 0,
-        "frequency_penalty": 1,
-        "response_format": None,
-    }
-
-    # Headers for the request, including the Authorization bearer token
-    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
-
-    # Make the POST request to the API
-    response = requests.post(url, json=payload, headers=headers)
-    response = json.loads(response.text)
-    print(response)
-    return [response["choices"][0]["message"]["content"], response["citations"]]
-
-
-def search_web(
-    query: str,
-    num_results: int = 5,
-    provider: str = "duckduckgo",
-    api_key=None,
-    **kwargs,
-) -> List[Dict[str, str]]:
-    """
-    Function Description:
-        This function searches the web for information based on a query.
-    Args:
-        query: The search query.
-    Keyword Args:
-        num_results: The number of search results to retrieve.
-        provider: The search engine provider to use ('google' or 'duckduckgo').
-    Returns:
-        A list of dictionaries with 'title', 'link', and 'content' keys.
-    """
-    results = []
-
-    # try:
-    if provider == "perplexity":
-        search_result = search_perplexity(query, api_key=api_key, **kwargs)
-        # print(search_result, type(search_result))
-        return search_result
-
-    if provider == "duckduckgo":
-        ddgs = DDGS()
-        try:
-            search_results = ddgs.text(query, max_results=num_results)
-            print(search_results, type(search_results))
-            urls = [r["href"] for r in search_results]
-            results = [
-                {"title": r["title"], "link": r["href"], "content": r["body"]}
-                for r in search_results
-            ]
-        except DuckDuckGoSearchException as e:
-            print("DuckDuckGo search failed: ", e)
-            urls = []
-            results = []
-
-    else:  # google
-        urls = list(search(query, num_results=num_results))
-        # google shit doesnt seem to be working anymore, apparently a lbock they made on browsers without js?
-        print("urls", urls)
-        print(provider)
-        for url in urls:
-            try:
-                # Fetch the webpage content
-                headers = {
-                    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
-                }
-                response = requests.get(url, headers=headers, timeout=5)
-                response.raise_for_status()
-
-                # Parse with BeautifulSoup
-                soup = BeautifulSoup(response.text, "html.parser")
-
-                # Get title and content
-                title = soup.title.string if soup.title else url
-
-                # Extract text content and clean it up
-                content = " ".join([p.get_text() for p in soup.find_all("p")])
-                content = " ".join(content.split())  # Clean up whitespace
-
-                results.append(
-                    {
-                        "title": title,
-                        "link": url,
-                        "content": (
-                            content[:500] + "..." if len(content) > 500 else content
-                        ),
-                    }
-                )
-
-            except Exception as e:
-                print(f"Error fetching {url}: {str(e)}")
-                continue
-
-    # except Exception as e:
-    #    print(f"Search error: {str(e)}")
-    content_str = "\n".join(
-        [r["content"] + "\n Citation: " + r["link"] + "\n\n\n" for r in results]
-    )
-    link_str = "\n".join([r["link"] + "\n" for r in results])
-    return [content_str, link_str]
-
-
-def rag_search(
-    query: str,
-    text_data: Union[Dict[str, str], str],
-    embedding_model: Any = None,
-    text_data_embedded: Optional[Dict[str, np.ndarray]] = None,
-    similarity_threshold: float = 0.3,
-    device="cpu",
-) -> List[str]:
-    """
-    Function Description:
-        This function retrieves lines from documents that are relevant to the query.
-    Args:
-        query: The query string.
-        text_data: A dictionary with file paths as keys and file contents as values.
-        embedding_model: The sentence embedding model.
-    Keyword Args:
-        text_data_embedded: A dictionary with file paths as keys and embedded file contents as values.
-        similarity_threshold: The similarity threshold for considering a line relevant.
-    Returns:
-        A list of relevant snippets.
-
-    """
-    if embedding_model is None:
-        try:
-            embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
-        except:
-            raise Exception(
-                "Please install the sentence-transformers library to use this function or provide an embedding transformer model."
-            )
-    results = []
-
-    # Compute the embedding of the query
-    query_embedding = embedding_model.encode(
-        query, convert_to_tensor=True, show_progress_bar=False
-    )
-    if isinstance(text_data, str):
-        # split at the sentence level
-        lines = text_data.split(".")
-        if not lines:
-            return results
-        # Compute embeddings for each line
-        if text_data_embedded is None:
-            line_embeddings = embedding_model.encode(lines, convert_to_tensor=True)
-        else:
-            line_embeddings = text_data_embedded
-        # Compute cosine similarities
-        cosine_scores = util.cos_sim(query_embedding, line_embeddings)[0].cpu().numpy()
-
-        # Find indices of lines above the similarity threshold
-        relevant_line_indices = np.where(cosine_scores >= similarity_threshold)[0]
-        # print("relevant_line_indices", cosine_scores)
-        # print(np.mean(cosine_scores))
-        # print(np.max(cosine_scores))
-
-        for idx in relevant_line_indices:
-            idx = int(idx)
-            # Get context lines (±10 lines)
-            start_idx = max(0, idx - 10)
-            end_idx = min(len(lines), idx + 11)  # +11 because end index is exclusive
-            snippet = ". ".join(lines[start_idx:end_idx])
-            results.append(snippet)
-
-    elif isinstance(text_data, dict):
-        for filename, content in text_data.items():
-            # Split content into lines
-            lines = content.split("\n")
-            if not lines:
-                continue
-            # Compute embeddings for each line
-            if text_data_embedded is None:
-                line_embeddings = embedding_model.encode(lines, convert_to_tensor=True)
-            else:
-                line_embeddings = text_data_embedded[filename]
-            # Compute cosine similarities
-            cosine_scores = (
-                util.cos_sim(query_embedding, line_embeddings)[0].cpu().numpy()
-            )
-
-            # Find indices of lines above the similarity threshold
-            ##print("most similar", np.max(cosine_scores))
-            ##print("most similar doc", lines[np.argmax(cosine_scores)])
-            relevant_line_indices = np.where(cosine_scores >= similarity_threshold)[0]
-            # print("relevant_line_indices", cosine_scores)
-            # print(np.mean(cosine_scores))
-            # print(np.max(cosine_scores))
-            for idx in relevant_line_indices:
-                idx = int(idx)  # Ensure idx is an integer
-                # Get context lines (±10 lines)
-                start_idx = max(0, idx - 10)
-                end_idx = min(
-                    len(lines), idx + 11
-                )  # +11 because end index is exclusive
-                snippet = "\n".join(lines[start_idx:end_idx])
-                results.append((filename, snippet))
-        # print("results", results)
-    return results