import os
import glob

ROAM_PATH = "/your/path/to/org-roam/files"
BANNED_FILES = {"my_journal",
                "personal_finances"}

def get_all_files_in_a_folder(full=False):
    # Use os.path.join() to safely join file paths
    path = os.path.join(ROAM_PATH, "**/*.org")

    # Use set operations to filter out unwanted files
    files = set(glob.glob(path, recursive=True))
    files = files.difference(set(f for f in files if any(word in f for word in BANNED_FILES)))

    if not full:
        # Use os.path.basename() to get just the file name
        files = [os.path.splitext(os.path.basename(file))[0] for file in files]

    return files

import orgparse
import pandas as pd
import re

def org_roam_nodes_to_dataframe(org_file):
    # Load the org file into an OrgData object
    org_data = orgparse.load(org_file)

    # Define a function to extract the title of a node
    def extract_title(node):
        if node.heading:
            # If the node has a heading, return it
            return node.heading
        else:
            # Otherwise, extract the title from the org file using a regular expression
            title_pattern = re.compile(r'^#\+title:\s*(.*)$', re.IGNORECASE)
            match = title_pattern.search(node.body)
            if match:
                return match.group(1)
            else:
                # If the title is not found, extract it from the first line of the body
                return re.sub(r"#\+title:",
				"",
				node.body.split("\n")[0], flags=re.IGNORECASE).strip()

    # Define a function to recursively extract the bodies of a node and its descendants
    def extract_node_nested_body(node):
        body = node.body
        for child in node.children:
            body += '\n' + child.level * "*" + " " + child.heading + "\n" + \
			extract_node_nested_body(child)
        return body.strip()

    # Define a function to recursively extract the bodies of a node
	# and its descendants when they are not other nodes
    def extract_node_nested_body_exclusive(node):
        body = node.body
        for child in node.children:
            if not child.properties.get('ID') and not child.properties.get('SEARCH'):
                body += '\n' + child.level * "*" + " " + child.heading + "\n" + \
				extract_node_nested_body_exclusive(child)
        return body.strip()

    # Define a function to build the hierarchy of a node
    def build_node_hierarchy(node):
        hierarchy = [extract_title(node)]
        parent = node.parent

        # while parent and parent != org_data[0]:
        while parent:
            hierarchy.append(extract_title(parent))
            parent = parent.parent
        return ' > '.join(reversed(hierarchy)).strip()

    # Define a function to convert a node to a dictionary
    def node_to_dict(node, file_name):
        node_dict = {
            'file_name': file_name,
            'node_id': node.properties.get('ID'),
            'node_title': extract_title(node),
            'node_hierarchy': build_node_hierarchy(node),
            'node_text': node.body,
            'node_text_nested': extract_node_nested_body(node),
            'node_text_nested_exclusive': extract_node_nested_body_exclusive(node),
        }

        return node_dict

    # Create a list of all org-roam nodes in the OrgData object
    nodes = [node_to_dict(node, org_file) for node in org_data[0][:] if node.properties.get('ID')]

    return pd.DataFrame(nodes)

org_files = get_all_files_in_a_folder(full=True)

roam_nodes_df = pd.concat([org_roam_nodes_to_dataframe(org_file) for org_file in org_files])

roam_nodes_df["text_to_encode"] = (
    roam_nodes_df["node_text_nested_exclusive"]
    .astype(str)
    .str.replace("#\+filetags:", "tags:")
    .str.replace("#\+title:", "title:")
    .str.replace("#\+STARTUP: inlineimages latexpreview", "")
)

roam_nodes_df["text_to_encode"] = (
    "[" + roam_nodes_df["node_hierarchy"] + "] " +
    roam_nodes_df["text_to_encode"].astype(str)
)

import tiktoken

TOKENS_PACK = 1000
COST_PER_TOKEN = 0.0004

encoding = tiktoken.get_encoding("cl100k_base")
tokens_per_text = roam_nodes_df["text_to_encode"].apply(encoding.encode).apply(len)

total_tokens = tokens_per_text.sum()
total_cost = total_tokens / TOKENS_PACK * COST_PER_TOKEN

print(total_cost)

from langchain.vectorstores import Chroma
from langchain.docstore.document import Document
from langchain.embeddings import OpenAIEmbeddings

persist_directory = "/path/to/the/db-org-roam-nodes-openai"

embedding = OpenAIEmbeddings(openai_api_key="<your-open-ai-key>")
vectordb = Chroma("langchain_store", embedding_function=embedding, persist_directory=persist_directory)

vectordb.add_texts(list(roam_nodes_df["node_text_to_encode"].values),
                   metadatas=[{"ID": roam_nodes_df.iloc[i]["node_id"],
                               "title": roam_nodes_df.iloc[i]["node_title"],
                               "hierarchy": roam_nodes_df.iloc[i]["node_hierarchy"],
                               "file_name": roam_nodes_df.iloc[i]["file_name"]} for i in range(len(roam_nodes_df))],
                   ids=list(roam_nodes_df["node_id"]))

vectordb.persist()

input_prompt = "A machine learning model is degrading. Its performance has been dropping since it was released. What are the possible reasons for it?"
retrieved_docs = vectordb.similarity_search_with_score(input_prompt, k=10)
retrieved_docs

import urllib.parse
import re

from http.server import BaseHTTPRequestHandler, HTTPServer
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma


class RequestHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.send_header('Content-type', 'text/plain')
        self.end_headers()

        # Input
        request_str = urllib.parse.unquote(self.path.split("/api/")[-1])
        print(request_str)

        # Retrieve docs
        retrieved_docs = vectordb.similarity_search_with_score(request_str, k=15)
        org_link_format = "[%.2f]: [[id:%s][%s]] \n %s"
        docs = [org_link_format % (score, doc.metadata["ID"],
		                           doc.metadata["title"].strip(),
	                                   doc.metadata["hierarchy"].strip())
                                 for doc, score in retrieved_docs]

        # Format the output
        response_str = f"#+title: Most similar nodes \n\n:QUERY:\n{request_str} \n:END:\n\n"
        for i, source in enumerate(docs):
            response_str += "* " + source + "\n"

        print(response_str)
        self.wfile.write(response_str.encode())

def run_server():
    server_address = ('', 8800)
    httpd = HTTPServer(server_address, RequestHandler)
    print(f'Server is running on port {server_address[1]}')

    httpd.serve_forever()

if __name__ == '__main__':
    persist_directory = "the/path/to/your/vector-db"
    embedding = OpenAIEmbeddings(openai_api_key="<your-openai-key>")

    vectordb = Chroma("langchain_store",
              embedding_function=embedding,
              persist_directory=persist_directory)
    run_server()

from langchain.text_splitter import CharacterTextSplitter

text_splitter = CharacterTextSplitter(chunk_size=300, chunk_overlap=0)

for index, row in data.iterrows():
    org_id = row["node_id"]
    title = row["node_title"]
    file_name = row["file_name"]
    node_hierarchy = row["node_hierarchy"]
    texts = text_splitter.split_text(row["node_text_nested_exclusive"])
    texts = ["[" + node_hierarchy + "] " + text for text in texts]
    metadatas = [
        {
            "source": f"{index}-{i}",
            "ID": org_id,
            "title": title,
            "hierarchy": node_hierarchy,
            "file_name": file_name,
        }
        for i in range(len(texts))
    ]
    ids = [f"{index}-{i}" for i in range(len(texts))]
    vectordb.add_texts(texts, metadatas=metadatas, ids=ids)

from typing import List, Tuple

def similarity_search_with_embeddings(
    vectordb,
    query: str,
    k: int = 4,
    filter: Optional[Dict[str, str]] = None,
    ,**kwargs: Any,
):
    """Run similarity search with Chroma and return labels, embeddings and scores of top k results.

    Args:
        query (str): Query text to search for.
        k (int): Number of results to return. Defaults to 4.
        filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

    Returns:
        Labels (List), Embeddings (List), Scores (List)
    """
    query_embedding = vectordb._embedding_function.embed_query(query)
    results = vectordb._collection.query(
        query_embedding,
        n_results=k,
        where=filter,
        include=["metadatas", "documents", "distances", "embeddings"], **kwargs
    )

    return [d["title"] for d in results["metadatas"][0]], results["embeddings"][0], results["distances"][0]

import matplotlib.pyplot as plt
import numpy as np

from sklearn.manifold import TSNE
from adjustText import adjust_text

def plot_vectors(vectors, labels, colors):
  shorten_labels = []
  for label in labels:
      if len(label) > 30:
          label = label[:30] + "..."
      shorten_labels.append(label)
  labels = shorten_labels
  # Project the vectors into two dimensions using t-SNE
  tsne = TSNE(n_components=2,
              perplexity=5)
  projected_vectors = tsne.fit_transform(np.array(vectors))

  # Create a colormap with different shades of purple
  cmap = plt.get_cmap("winter")
  plt.figure(figsize=(12, 12))
  # Plot the projected vectors with different shades of purple
  plt.scatter(projected_vectors[:, 0], projected_vectors[:, 1], c=colors, cmap=cmap)

  # Add labels to the plot using ajustText
  texts = []
  for label, x, y in zip(labels, projected_vectors[:, 0], projected_vectors[:, 1]):
    text = plt.annotate(label, xy=(x, y), fontsize=15)
    texts.append(text)
  adjust_text(texts)

  plt.savefig("path/to/the/output.jpg")