ziya 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

app/agents/agent.py

@@ -1,8 +1,7 @@
 import os
-import botocore
-
-from typing import Generator, List, Tuple, Set, Union
+from typing import List, Tuple, Set, Union
 
+import botocore
 import tiktoken
 from langchain.agents import AgentExecutor
 from langchain.agents.format_scratchpad import format_xml
@@ -12,12 +11,12 @@ from langchain_core.messages import AIMessage, HumanMessage
 from langchain_core.pydantic_v1 import BaseModel, Field
 
 from app.agents.prompts import conversational_prompt, parse_output
-from app.utils.directory_util import get_ignored_patterns, get_complete_file_list
 from app.utils.logging_utils import logger
 from app.utils.print_tree_util import print_file_tree
 
 
 def _format_chat_history(chat_history: List[Tuple[str, str]]) -> List[Union[HumanMessage, AIMessage]]:
+    logger.info("Formatting chat history")
     buffer = []
     for human, ai in chat_history:
         buffer.append(HumanMessage(content=human))
@@ -46,32 +45,20 @@ model = ChatBedrock(
 )
 
 
-def get_combined_document_contents() -> str:
-    user_codebase_dir: str = os.environ["ZIYA_USER_CODEBASE_DIR"]
-    print(f"Reading user's current codebase: {user_codebase_dir}")
-
+def get_combined_docs_from_files(files) -> str:
     combined_contents: str = ""
-    ignored_patterns: List[str] = get_ignored_patterns(user_codebase_dir)
-    included_relative_dirs = [pattern.strip() for pattern in os.environ.get("ZIYA_ADDITIONAL_INCLUDE_DIRS", "").split(',')]
-    all_files: List[str] = get_complete_file_list(user_codebase_dir, ignored_patterns, included_relative_dirs)
-
-    print_file_tree(all_files)
-
-    seen_dirs: Set[str] = set()
-    for file_path in all_files:
+    print_file_tree(files)
+    user_codebase_dir: str = os.environ["ZIYA_USER_CODEBASE_DIR"]
+    for file_path in files:
         try:
-            docs = TextLoader(file_path).load()
+            full_file_path = os.path.join(user_codebase_dir, file_path)
+            docs = TextLoader(full_file_path).load()
             for doc in docs:
                 combined_contents += f"File: {file_path}\n{doc.page_content}\n\n"
-            dir_path = os.path.dirname(file_path)
-            if dir_path not in seen_dirs:
-                seen_dirs.add(dir_path)
+
         except Exception as e:
-            print(f"Skipping file {file_path} due to error: {e}")
+            print(f"Skipping file {full_file_path} due to error: {e}")
 
-    print("--------------------------------------------------------")
-    print(f"Ignoring following paths based on gitIgnore and other defaults: {ignored_patterns}")
-    print("--------------------------------------------------------")
     print(f"Codebase word count: {len(combined_contents.split()):,}")
     token_count = len(tiktoken.get_encoding("cl100k_base").encode(combined_contents))
     print(f"Codebase token count: {token_count:,}")
@@ -79,24 +66,57 @@ def get_combined_document_contents() -> str:
     print("--------------------------------------------------------")
     return combined_contents
 
-def get_child_paths(directory: str) -> Generator[str, None, None]:
-    for entry in os.listdir(directory):
-        child_path = os.path.join(directory, entry)
-        yield child_path
-
-
-prompt = conversational_prompt.partial(
-    codebase=get_combined_document_contents,
-)
+# def get_combined_document_contents_from_included_dirs(included_relative_dirs) -> str:
+#     user_codebase_dir: str = os.environ["ZIYA_USER_CODEBASE_DIR"]
+#     print(f"Reading user's current codebase: {user_codebase_dir}")
+#
+#     combined_contents: str = ""
+#     ignored_patterns: List[str] = get_ignored_patterns(user_codebase_dir)
+#     all_files: List[str] = get_complete_file_list(user_codebase_dir, ignored_patterns, included_relative_dirs)
+#     print("all_files")
+#     print(all_files)
+#
+#     print_file_tree(all_files)
+#
+#     seen_dirs: Set[str] = set()
+#     for file_path in all_files:
+#         try:
+#             docs = TextLoader(file_path).load()
+#             for doc in docs:
+#                 combined_contents += f"File: {file_path}\n{doc.page_content}\n\n"
+#             dir_path = os.path.dirname(file_path)
+#             if dir_path not in seen_dirs:
+#                 seen_dirs.add(dir_path)
+#         except Exception as e:
+#             print(f"Skipping file {file_path} due to error: {e}")
+#
+#     print("--------------------------------------------------------")
+#     print(f"Ignoring following paths based on gitIgnore and other defaults: {ignored_patterns}")
+#     print("--------------------------------------------------------")
+#     print(f"Codebase word count: {len(combined_contents.split()):,}")
+#     token_count = len(tiktoken.get_encoding("cl100k_base").encode(combined_contents))
+#     print(f"Codebase token count: {token_count:,}")
+#     print(f"Max Claude Token limit: 200,000")
+#     print("--------------------------------------------------------")
+#     return combined_contents
+
+# prompt = conversational_prompt.partial(
+#     # codebase=get_combined_document_contents()
+# )
 llm_with_stop = model.bind(stop=["</tool_input>"])
 
+def extract_codebase(x):
+    # return get_combined_docs_from_files(x["config"].get("include_dirs", ['./']))
+    return get_combined_docs_from_files(x["config"].get("files", []))
+
 agent = (
         {
+            "codebase": lambda x: extract_codebase(x),
             "question": lambda x: x["question"],
             "agent_scratchpad": lambda x: format_xml(x["intermediate_steps"]),
             "chat_history": lambda x: _format_chat_history(x["chat_history"]),
         }
-        | prompt
+        | conversational_prompt
         | llm_with_stop
         | parse_output
 )
@@ -104,6 +124,7 @@ agent = (
 
 class AgentInput(BaseModel):
     question: str
+    config: dict = Field({})
     chat_history: List[Tuple[str, str]] = Field(..., extra={"widget": {"type": "chat"}})
 
 

app/agents/prompts.py

@@ -2,7 +2,8 @@ from langchain_core.agents import AgentFinish
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 # import pydevd_pycharm
 
-template = """You are an excellent coder. Help the user with their coding tasks. You are given the entire codebase
+template = """
+You are an excellent coder. Help the user with their coding tasks. You are given the entire codebase
  of the user in your context. It is in the format like below where first line has the File path and then the content follows. 
 
 File: <filepath>
@@ -16,6 +17,7 @@ Now below is the current codebase of the user:
 conversational_prompt = ChatPromptTemplate.from_messages(
     [
         ("system", template),
+        # ("system", "You are a helpful AI bot. Your name is {name}."),
         MessagesPlaceholder(variable_name="chat_history"),
         ("user", "{question}"),
         ("ai", "{agent_scratchpad}"),

app/main.py

@@ -15,8 +15,6 @@ def parse_arguments():
     parser = argparse.ArgumentParser(description="Run with custom options")
     parser.add_argument("--exclude", default=[], type=lambda x: x.split(','),
                         help="List of files or directories to exclude (e.g., --exclude 'tst,build,*.py')")
-    parser.add_argument("--include", default=[], type=lambda x: x.split(','),
-                        help="List of directories to include (e.g., --include 'src,static'). Only directories for now")
     parser.add_argument("--profile", type=str, default=None,
                         help="AWS profile to use (e.g., --profile ziya)")
     parser.add_argument("--model", type=str, choices=["sonnet", "haiku", "opus"], default="sonnet",
@@ -34,9 +32,6 @@ def setup_environment(args):
     additional_excluded_dirs = ','.join(args.exclude)
     os.environ["ZIYA_ADDITIONAL_EXCLUDE_DIRS"] = additional_excluded_dirs
 
-    additional_included_dirs = ','.join(args.include)
-    os.environ["ZIYA_ADDITIONAL_INCLUDE_DIRS"] = additional_included_dirs
-
     if args.profile:
         os.environ["ZIYA_AWS_PROFILE"] = args.profile
     if args.model:

app/server.py

@@ -1,22 +1,88 @@
+import os
+from typing import Dict, Any, List, Tuple
+
+import tiktoken
 from fastapi import FastAPI, Request
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from langserve import add_routes
 from app.agents.agent import agent_executor
+from fastapi.responses import FileResponse
 
+# import pydevd_pycharm
 import uvicorn
 
+from app.utils.directory_util import get_ignored_patterns
+from app.utils.gitignore_parser import parse_gitignore_patterns
 
 app = FastAPI()
-app.mount("/static", StaticFiles(directory="../static"), name="static")
+app.mount("/static", StaticFiles(directory="../templates/static"), name="static")
 templates = Jinja2Templates(directory="../templates")
 
 # Add a route for the frontend
 add_routes(app, agent_executor, disabled_endpoints=["playground"], path="/ziya")
 
+
 @app.get("/")
 async def root(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})
 
+
+@app.get("/favicon.ico", include_in_schema=False)
+async def favicon():
+    return FileResponse('../templates/favicon.ico')
+
+
+def get_folder_structure(directory: str, ignored_patterns: List[Tuple[str, str]], max_depth: int = 99) -> Dict[str, Any]:
+    should_ignore_fn = parse_gitignore_patterns(ignored_patterns)
+
+    def count_tokens(file_path: str) -> int:
+        try:
+            with open(file_path, 'r') as file:
+                content = file.read()
+                return len(tiktoken.get_encoding("cl100k_base").encode(content))
+        except Exception as e:
+            print(f"Skipping file {file_path} due to error: {e}")
+            return 0
+
+    def get_structure(current_dir: str, current_depth: int):
+        if current_depth > max_depth:
+            return None
+
+        current_structure = {}
+        for entry in os.listdir(current_dir):
+            if entry.startswith('.'):  # Skip hidden files/folders
+                continue
+            entry_path = os.path.join(current_dir, entry)
+            if os.path.isdir(entry_path):
+                if not should_ignore_fn(entry_path):
+                    sub_structure = get_structure(entry_path, current_depth + 1)
+                    if sub_structure is not None:
+                        token_count = sum(sub_structure[key]['token_count'] for key in sub_structure)
+                        current_structure[entry] = {'token_count': token_count, 'children': sub_structure}
+            else:
+                if not should_ignore_fn(entry_path):
+                    token_count = count_tokens(entry_path)
+                    current_structure[entry] = {'token_count': token_count}
+
+        return current_structure
+
+    folder_structure = get_structure(directory, 1)
+    return folder_structure
+
+
+@app.get("/api/folders")
+async def get_folders():
+    # pydevd_pycharm.settrace('localhost', port=59939, stdoutToServer=True, stderrToServer=True)
+    user_codebase_dir = os.environ["ZIYA_USER_CODEBASE_DIR"]
+    ignored_patterns: List[Tuple[str, str]] = get_ignored_patterns(user_codebase_dir)
+    return get_folder_structure(user_codebase_dir, ignored_patterns)
+
+
+@app.get('/api/default-included-folders')
+def get_default_included_folders():
+    return {'defaultIncludedFolders': []}
+
+
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)

app/utils/directory_util.py

@@ -1,6 +1,6 @@
 import glob
 import os
-from typing import List, Set, Tuple
+from typing import List, Set, Tuple, Dict
 
 from app.utils.gitignore_parser import parse_gitignore_patterns
 
@@ -41,20 +41,25 @@ def get_ignored_patterns(directory: str) -> List[Tuple[str, str]]:
     return ignored_patterns
 
 
-def get_complete_file_list(user_codebase_dir: str, ignored_patterns: List[str], included_relative_dirs: List[str]) -> List[str]:
-    should_ignore = parse_gitignore_patterns(ignored_patterns, base_dir=user_codebase_dir)
-    file_set: Set[str] = set()
+def get_complete_file_list(user_codebase_dir: str, ignored_patterns: List[str], included_relative_dirs: List[str]) -> Dict[str, Dict]:
+    should_ignore_fn = parse_gitignore_patterns(ignored_patterns)
+    file_dict: Dict[str, Dict] = {}
     for pattern in included_relative_dirs:
         for root, dirs, files in os.walk(os.path.normpath(os.path.join(user_codebase_dir, pattern))):
-            # Filter out ignored directories
-            dirs[:] = [d for d in dirs if not should_ignore(os.path.join(root, d))]
+            # Filter out ignored directories and hidden directories
+            dirs[:] = [d for d in dirs if not should_ignore_fn(os.path.join(root, d)) and not d.startswith('.')]
 
             for file in files:
                 file_path = os.path.join(root, file)
-                if not should_ignore(file_path) and not is_image_file(file_path):
-                    file_set.add(file_path)
+                if not should_ignore_fn(file_path) and not is_image_file(file_path) and not file.startswith('.'):
+                    file_dict[file_path] = {}
 
-    return list(file_set)
+            for dir in dirs:
+                dir_path = os.path.join(root, dir)
+                if not should_ignore_fn(dir_path):
+                    file_dict[dir_path] = {}
+
+    return file_dict
 
 def is_image_file(file_path: str) -> bool:
     image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.svg', '.ico']

app/utils/gitignore_parser.py

@@ -14,9 +14,7 @@ def handle_negation(file_path, rules: Reversible["IgnoreRule"]):
     return False
 
 
-def parse_gitignore_patterns(patterns: List[Tuple[str, str]], base_dir=None):
-    if base_dir is None:
-        base_dir = os.getcwd()
+def parse_gitignore_patterns(patterns: List[Tuple[str, str]]):
     rules = []
     for counter, (pattern, directory) in enumerate(patterns, start=1):
         rule = rule_from_pattern(pattern, base_path=_normalize_path(directory),

pyproject.toml

@@ -1,10 +1,10 @@
 [tool.poetry]
 name = "ziya"
-version = "0.1.7"
+version = "0.1.9"
 description = ""
 authors = ["Vishnu Krishnaprasad <vishnukool@gmail.com>"]
 readme = "README.md"
-include = ["static/**/*", "templates/**/*", "pyproject.toml"]
+include = ["templates/**/*", "pyproject.toml" ,"scripts.py"]
 packages = [
     { include = "app" },
 ]
@@ -21,8 +21,6 @@ langchain = "^0.1"
 langchainhub = ">=0.1.15"
 langchain-anthropic = "^0.1.4"
 langchain-cli = ">=0.0.15"
-
-[tool.poetry.group.dev.dependencies]
 pydevd-pycharm = "^241.15989.57"
 
 [build-system]
@@ -30,4 +28,8 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry.scripts]
-ziya = 'app.main:main'
+ziya = 'scripts:ziya'
+dev = 'scripts:dev'
+finstall = "scripts:frontend_install"
+fstart = "scripts:frontend_start"
+fbuild = "scripts:frontend_build"

scripts.py

@@ -0,0 +1,42 @@
+import threading
+import subprocess
+import signal
+import sys
+
+from app.main import main
+
+
+def frontend_start():
+    subprocess.run(["npm", "run", "start"], cwd="frontend")
+
+
+def frontend_install():
+    subprocess.run(["npm", "install"], cwd="frontend")
+
+
+def frontend_build():
+    subprocess.run(["npm", "run", "build"], cwd="frontend")
+
+
+def ziya():
+    # frontend_build()
+    main()
+
+
+def signal_handler(sig, frame):
+    print('Interrupt received, shutting down...')
+    sys.exit(0)
+
+
+def dev():
+    frontend_thread = threading.Thread(target=frontend_start)
+    signal.signal(signal.SIGINT, signal_handler)
+
+    frontend_thread.start()
+    try:
+        print("Camme to main")
+        main()
+    except KeyboardInterrupt:
+        print("Main process interrupted.")
+    finally:
+        frontend_thread.join()

templates/asset-manifest.json

@@ -0,0 +1,21 @@
+{
+  "files": {
+    "main.css": "/static/css/main.85c755f1.css",
+    "main.js": "/static/js/main.d47d0f10.js",
+    "static/media/fa-solid-900.ttf": "/static/media/fa-solid-900.bacd5de623fb563b961a.ttf",
+    "static/media/fa-brands-400.ttf": "/static/media/fa-brands-400.60127e352b7a11f7f1bc.ttf",
+    "static/media/fa-solid-900.woff2": "/static/media/fa-solid-900.4d986b00ff9ca3828fbd.woff2",
+    "static/media/fa-brands-400.woff2": "/static/media/fa-brands-400.455ea818179b4def0c43.woff2",
+    "static/media/fa-regular-400.ttf": "/static/media/fa-regular-400.eb91f7b948a42799f678.ttf",
+    "static/media/fa-regular-400.woff2": "/static/media/fa-regular-400.21cb8f55d8e0c5b89751.woff2",
+    "static/media/fa-v4compatibility.ttf": "/static/media/fa-v4compatibility.c8e090db312b0bea2aa2.ttf",
+    "static/media/fa-v4compatibility.woff2": "/static/media/fa-v4compatibility.cf7f5903d06b79ad60f1.woff2",
+    "index.html": "/index.html",
+    "main.85c755f1.css.map": "/static/css/main.85c755f1.css.map",
+    "main.d47d0f10.js.map": "/static/js/main.d47d0f10.js.map"
+  },
+  "entrypoints": [
+    "static/css/main.85c755f1.css",
+    "static/js/main.d47d0f10.js"
+  ]
+}

templates/index.html

@@ -1,19 +1 @@
-<!DOCTYPE html>
-<html lang="en">
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>Ziya - Code Assistant</title>
-    <link rel="stylesheet" href="/static/ziya.css">
-    <link rel="icon" href="/static/favicon.ico" type="image/x-icon">
-</head>
-<body>
-<div id="root"/>
-<script src="https://cdn.jsdelivr.net/npm/@babel/standalone@7/babel.min.js"></script>
-<script src="https://unpkg.com/react@17/umd/react.production.min.js"></script>
-<script src="https://unpkg.com/react-dom@17/umd/react-dom.production.min.js"></script>
-<script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
-<script type="text/babel" src="/static/sendPayload.js"></script>
-<script type="text/babel" src="/static/app.js"></script>
-</body>
-</html>
+<!doctype html><html lang="en"><head><meta charset="UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1"><title>Ziya - Code Assistant</title><link rel="icon" href="/favicon.ico" type="image/x-icon"><script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script><script defer="defer" src="/static/js/main.d47d0f10.js"></script><link href="/static/css/main.85c755f1.css" rel="stylesheet"></head><body><div id="root"></div></body></html>