feat: make incognito work

.vscode/launch.json

@@ -0,0 +1,19 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: FastAPI",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "uvicorn",
+            "args": [
+                "server:app",
+                "--reload"
+            ],
+            "jinja": true
+        }
+    ]
+}

README.md

@@ -18,11 +18,11 @@ LlamaFS runs in two "modes" - as a batch job (batch mode), and an interactive da
 
 In batch mode, you can send a directory to LlamaFS, and it will return a suggested file structure and organize your files.
 
-In watch mode, LlamaFS starts a daemon that watches your directory. It intercepts all filesystem operations and uses your most recent edits to proactively learn how you rename file. For example, if you create a folder for your 2023 tax documents, and start moving 1-3 files in it, LlamaFS will automatically create and move the files for you!
+In watch mode, LlamaFS starts a daemon that watches your directory. It intercepts all filesystem operations and uses your most recent edits to proactively learn how you rename file. For example, if you create a folder for your 2023 tax documents, and start moving 1-3 files in it, LlamaFS will automatically create and move the files for you! (watch mode defaults to sending files to groq if you have the environment variable "GROQ_API_KEY" set, otherwise through ollama)
 
 Uh... Sending all my personal files to an API provider?! No thank you!
 
-It also has a toggle for "incognito mode," allowing you route every request through Ollama instead of Groq. Since they use the same Llama 3 model, the perform identically.
+BREAKING CHANGE: Now by default, llama-fs uses "incognito mode" (if you have not configured an environment key for "GROQ_API_KEY") allowing you route every request through Ollama instead of Groq. Since they use the same Llama 3 model, the perform identically. To use a different model, set the environment variable "MODEL" to a string which litellm can use as a model like "ollama/llama3" or "groq/llama3-70b-8192".
 
 ## How we built it
 
@@ -42,7 +42,7 @@ We built LlamaFS on a Python backend, leveraging the Llama3 model through Groq f
 ### Prerequisites
 
 Before installing, ensure you have the following requirements:
-- Python 3.10 or higher
+- Python 3.9 or higher
 - pip (Python package installer)
 
 ### Installing
@@ -63,11 +63,11 @@ To install the project, follow these steps:
    pip install -r requirements.txt
    ```
 
-4. (Optional) Install moondream if you
-   want to use the incognito mode
+4. Install ollama and pull the model moondream if you want to recognize images
     ```bash
     ollama pull moondream
     ```
+   We highly recommend pulling an additional model like llama3 for local ai inference on text files. You can control which ollama model is used by setting the "MODEL" environment variable to a litellm compatible model string.
 
 ## Usage
 

main.py

@@ -2,7 +2,6 @@
 import json
 import argparse
 import pathlib
-from groq import Groq
 from llama_index.core import SimpleDirectoryReader
 import colorama
 import pathlib
@@ -24,10 +23,6 @@
 @click.argument("dst_path", type=click.Path())
 @click.option("--auto-yes", is_flag=True, help="Automatically say yes to all prompts")
 def main(src_path, dst_path, auto_yes=False):
-    os.environ["GROQ_API_KEY"] = (
-        "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"
-    )
-
     summaries = asyncio.run(get_dir_summaries(src_path))
 
     # Get file tree

requirements.txt

@@ -2,7 +2,6 @@ ollama
 chromadb
 llama-index
 litellm
-groq
 docx2txt
 colorama
 termcolor

server.py

@@ -10,14 +10,12 @@
 
 
 import colorama
-import ollama
 import threading
 from asciitree import LeftAligned
 from asciitree.drawing import BOX_LIGHT, BoxStyle
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
-from groq import Groq
 from llama_index.core import SimpleDirectoryReader
 from pydantic import BaseModel
 from termcolor import colored
@@ -28,13 +26,10 @@
 from src.watch_utils import Handler
 from src.watch_utils import create_file_tree as create_watch_file_tree
 
-os.environ["GROQ_API_KEY"] = "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"
-
-
 class Request(BaseModel):
     path: Optional[str] = None
     instruction: Optional[str] = None
-    incognito: Optional[bool] = False
+    incognito: Optional[bool] = True
 
 
 class CommitRequest(BaseModel):
@@ -71,9 +66,9 @@ async def batch(request: Request):
         raise HTTPException(
             status_code=400, detail="Path does not exist in filesystem")
 
-    summaries = await get_dir_summaries(path)
+    summaries = await get_dir_summaries(path, incognito=request.incognito)
     # Get file tree
-    files = create_file_tree(summaries)
+    files = create_file_tree(summaries, incognito=request.incognito)
 
     # Recursively create dictionary from file paths
     tree = {}
@@ -107,7 +102,7 @@ async def watch(request: Request):
 
     observer = Observer()
     event_handler = Handler(path, create_watch_file_tree, response_queue)
-    await event_handler.set_summaries()
+    await event_handler.set_summaries(incognito=request.incognito)
     observer.schedule(event_handler, path, recursive=True)
     observer.start()
 

src/loader.py

@@ -1,26 +1,30 @@
 import asyncio
+import http
+import http.server
 import json
 import os
 from collections import defaultdict
 
 import agentops
 import colorama
-import ollama
 import weave
-from groq import AsyncGroq, Groq
+import litellm
+import ollama
 from llama_index.core import Document, SimpleDirectoryReader
 from llama_index.core.schema import ImageDocument
 from llama_index.core.node_parser import TokenTextSplitter
 from termcolor import colored
 
+from src import select_model
+
 
 # @weave.op()
 # @agentops.record_function("summarize")
-async def get_dir_summaries(path: str):
+async def get_dir_summaries(path: str, incognito=True):
     doc_dicts = load_documents(path)
     # metadata = process_metadata(doc_dicts)
 
-    summaries = await get_summaries(doc_dicts)
+    summaries = await get_summaries(doc_dicts, incognito=incognito)
 
     # Convert path to relative path
     for summary in summaries:
@@ -90,7 +94,7 @@ def process_metadata(doc_dicts):
     return metadata_list
 
 
-async def summarize_document(doc, client):
+async def summarize_document(doc, incognito = True):
     PROMPT = """
 You will be provided with the contents of a file along with its metadata. Provide a summary of the contents. The purpose of the summary is to organize files based on their content. To this end provide a concise but informative summary. Make the summary as specific to the file as possible.
 
@@ -108,12 +112,12 @@ async def summarize_document(doc, client):
     attempt = 0
     while attempt < max_retries:
         try:
-            chat_completion = await client.chat.completions.create(
+            chat_completion = litellm.completion(
                 messages=[
                     {"role": "system", "content": PROMPT},
                     {"role": "user", "content": json.dumps(doc)},
                 ],
-                model="llama3-70b-8192",
+                model=select_model(incognito),
                 response_format={"type": "json_object"},
                 temperature=0,
             )
@@ -135,7 +139,7 @@ async def summarize_document(doc, client):
     return summary
 
 
-async def summarize_image_document(doc: ImageDocument, client):
+async def summarize_image_document(doc: ImageDocument):
     PROMPT = """
 You will be provided with an image along with its metadata. Provide a summary of the image contents. The purpose of the summary is to organize files based on their content. To this end provide a concise but informative summary. Make the summary as specific to the file as possible.
 
@@ -152,7 +156,6 @@ async def summarize_image_document(doc: ImageDocument, client):
     client = ollama.AsyncClient()
     chat_completion = await client.chat(
         messages=[
-            # {"role": "system", "content": "Respond with one short sentence."},
             {
                 "role": "user",
                 "content": "Summarize the contents of this image.",
@@ -162,7 +165,6 @@ async def summarize_image_document(doc: ImageDocument, client):
         model="moondream",
         # format="json",
         # stream=True,
-        options={"num_predict": 128},
     )
 
     summary = {
@@ -176,21 +178,18 @@ async def summarize_image_document(doc: ImageDocument, client):
     return summary
 
 
-async def dispatch_summarize_document(doc, client):
+async def dispatch_summarize_document(doc, incognito=True):
     if isinstance(doc, ImageDocument):
-        return await summarize_image_document(doc, client)
+        return await summarize_image_document(doc)
     elif isinstance(doc, Document):
-        return await summarize_document({"content": doc.text, **doc.metadata}, client)
+        return await summarize_document({"content": doc.text, **doc.metadata}, incognito=incognito)
     else:
         raise ValueError("Document type not supported")
 
 
-async def get_summaries(documents):
-    client = AsyncGroq(
-        api_key=os.environ.get("GROQ_API_KEY"),
-    )
+async def get_summaries(documents, incognito=True):
     summaries = await asyncio.gather(
-        *[dispatch_summarize_document(doc, client) for doc in documents]
+        *[dispatch_summarize_document(doc, incognito=incognito) for doc in documents]
     )
     return summaries
 
@@ -219,88 +218,12 @@ def merge_summary_documents(summaries, metadata_list):
 ################################################################################################
 
 
-def get_file_summary(path: str):
-    client = Groq(
-        api_key=os.environ.get("GROQ_API_KEY"),
-    )
+def get_file_summary(path: str, incognito=True):
     reader = SimpleDirectoryReader(input_files=[path]).iter_data()
 
     docs = next(reader)
     splitter = TokenTextSplitter(chunk_size=6144)
     text = splitter.split_text("\n".join([d.text for d in docs]))[0]
     doc = Document(text=text, metadata=docs[0].metadata)
-    summary = dispatch_summarize_document_sync(doc, client)
-    return summary
-
-
-def dispatch_summarize_document_sync(doc, client):
-    if isinstance(doc, ImageDocument):
-        return summarize_image_document_sync(doc, client)
-    elif isinstance(doc, Document):
-        return summarize_document_sync({"content": doc.text, **doc.metadata}, client)
-    else:
-        raise ValueError("Document type not supported")
-
-
-def summarize_document_sync(doc, client):
-    PROMPT = """
-You will be provided with the contents of a file along with its metadata. Provide a summary of the contents. The purpose of the summary is to organize files based on their content. To this end provide a concise but informative summary. Make the summary as specific to the file as possible.
-
-Write your response a JSON object with the following schema:
-
-```json 
-{
-    "file_path": "path to the file including name",
-    "summary": "summary of the content"
-}
-```
-""".strip()
-
-    chat_completion = client.chat.completions.create(
-        messages=[
-            {"role": "system", "content": PROMPT},
-            {"role": "user", "content": json.dumps(doc)},
-        ],
-        model="llama3-70b-8192",
-        response_format={"type": "json_object"},
-        temperature=0,
-    )
-    summary = json.loads(chat_completion.choices[0].message.content)
-
-    try:
-        print(colored(summary["file_path"], "green"))  # Print the filename in green
-        print(summary["summary"])  # Print the summary of the contents
-        print("-" * 80 + "\n")  # Print a separator line with spacing for readability
-    except KeyError as e:
-        print(e)
-        print(summary)
-
-    return summary
-
-
-def summarize_image_document_sync(doc: ImageDocument, client):
-    client = ollama.Client()
-    chat_completion = client.chat(
-        messages=[
-            {
-                "role": "user",
-                "content": "Summarize the contents of this image.",
-                "images": [doc.image_path],
-            },
-        ],
-        model="moondream",
-        # format="json",
-        # stream=True,
-        options={"num_predict": 128},
-    )
-
-    summary = {
-        "file_path": doc.image_path,
-        "summary": chat_completion["message"]["content"],
-    }
-
-    print(colored(summary["file_path"], "green"))  # Print the filename in green
-    print(summary["summary"])  # Print the summary of the contents
-    print("-" * 80 + "\n")  # Print a separator line with spacing for readability
-
+    summary = dispatch_summarize_document(doc, incognito=incognito)
     return summary

src/select_model.py

@@ -0,0 +1,16 @@
+from os import environ
+from litellm import validate_environment
+import warnings
+
+def select_model(incognito=True):
+    model = "groq/llama3-70b-8192" if environ.get("GROQ_API_KEY") and incognito is False else environ.get("MODEL", "ollama/llama3")
+    litellm_validation = validate_environment(model)
+    if litellm_validation.get('keys_in_environment') is False:
+        raise EnvironmentError({
+            "errno": 1, 
+            "strerr": f"missing environment variables for model {model}", 
+            "missing_keys": ','.join(litellm_validation.get("missing_keys"))
+        })
+    if "ollama" not in model:
+        warnings.warn(f"sending the contents of your files to {model}!")
+    return model

src/tree_generator.py

@@ -1,7 +1,9 @@
-from groq import Groq
+import litellm
 import json
 import os
 
+from src import select_model
+
 FILE_PROMPT = """
 You will be provided with list of source files and a summary of their contents. For each file, propose a new path and filename, using a directory structure that optimally organizes the files using known conventions and best practices.
 Follow good naming conventions. Here are a few guidelines
@@ -27,15 +29,14 @@
 """.strip()
 
 
-def create_file_tree(summaries: list):
-    client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
-    chat_completion = client.chat.completions.create(
+def create_file_tree(summaries: list, incognito=True):
+    chat_completion = litellm.completion(
         messages=[
             {"role": "system", "content": FILE_PROMPT},
             {"role": "user", "content": json.dumps(summaries)},
         ],
-        model="llama3-70b-8192",
-        response_format={"type": "json_object"},  # Uncomment if needed
+        model=select_model(incognito),
+        response_format={"type": "json_object"},
         temperature=0,
     )