ziya 0.1.47__py3-none-any.whl → 0.1.48__py3-none-any.whl

app/agents/agent.py

@@ -1,23 +1,25 @@
+import json
 import os
 import os.path
-from typing import List, Tuple, Set, Union
+from typing import List, Tuple, Union
 
-import json
 import botocore
 import tiktoken
 from langchain.agents import AgentExecutor
 from langchain.agents.format_scratchpad import format_xml
+from langchain.chat_models.base import BaseChatModel
 from langchain_aws import ChatBedrock
 from langchain_community.document_loaders import TextLoader
 from langchain_core.agents import AgentFinish
 from langchain_core.messages import AIMessage, HumanMessage
-from pydantic import BaseModel, Field
+from langchain_google_genai import ChatGoogleGenerativeAI
 
+from app.utils.llm_constants import MODEL_MAPPING, AgentInput, SAMPLE_QUESTION, GEMINI_PREFIX, GOOGLE_API_KEY
 from app.agents.prompts import conversational_prompt
-from app.utils.sanitizer_util import clean_backtick_sequences
-
 from app.utils.logging_utils import logger
 from app.utils.print_tree_util import print_file_tree
+from app.utils.sanitizer_util import clean_backtick_sequences
+
 
 def clean_chat_history(chat_history: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
     """Clean chat history by removing invalid messages and normalizing content."""
@@ -44,29 +46,36 @@ def parse_output(message):
     text = clean_backtick_sequences(message.content)
     return AgentFinish(return_values={"output": text}, log=text)
 
-aws_profile = os.environ.get("ZIYA_AWS_PROFILE")
-if aws_profile:
-    logger.info(f"Using AWS Profile: {aws_profile}")
-else:
-    logger.info("No AWS profile specified via --aws-profile flag, using default credentials")
-model_id = {
-    "sonnet3.5": "us.anthropic.claude-3-5-sonnet-20240620-v1:0",
-    "sonnet3.5-v2": "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
-    "opus": "us.anthropic.claude-3-opus-20240229-v1:0",
-    "sonnet": "us.anthropic.claude-3-sonnet-20240229-v1:0",
-    "haiku": "us.anthropic.claude-3-haiku-20240307-v1:0",
-}[os.environ.get("ZIYA_AWS_MODEL", "sonnet3.5-v2")]
-logger.info(f"Using Claude Model: {model_id}")
-
-model = ChatBedrock(
-    model_id=model_id,
-    model_kwargs={"max_tokens": 4096, "temperature": 0.3, "top_k": 15},
-    credentials_profile_name=aws_profile if aws_profile else None,
-    config=botocore.config.Config(
-        read_timeout=900
-    )
-)
 
+def get_chat_model() -> BaseChatModel:
+    model_name = os.environ.get("ZIYA_AWS_MODEL", "sonnet3.5-v2")
+    
+    logger.info(f"Using model name: {model_name}")
+
+    model_mapped_name = MODEL_MAPPING[model_name]
+
+    if model_name.startswith(GEMINI_PREFIX):
+        api_key = os.environ.get(GOOGLE_API_KEY)
+        if not api_key:
+            raise ValueError("%s environment variable is required for Gemini model" % GOOGLE_API_KEY)
+
+        return ChatGoogleGenerativeAI(
+            model=model_mapped_name,
+            temperature=0.2,
+            max_output_tokens=4096,
+            top_k=15,
+            google_api_key=api_key,
+            timeout=None,
+            verbose=True,
+        )
+    else:
+        aws_profile = os.environ.get("ZIYA_AWS_PROFILE", None)
+        return ChatBedrock(
+            model_id=MODEL_MAPPING[model_name],
+            model_kwargs={"max_tokens": 4096, "temperature": 0.3, "top_k": 15},
+            credentials_profile_name=aws_profile,
+            config=botocore.config.Config(read_timeout=900),
+        )
 
 def get_combined_docs_from_files(files) -> str:
     combined_contents: str = ""
@@ -86,17 +95,17 @@ def get_combined_docs_from_files(files) -> str:
     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("--------------------------------------------------------")
+    print("-" * 120)
     return combined_contents
 
-
-llm_with_stop = model.bind(stop=["</tool_input>"])
-
 def extract_codebase(x):
     logger.debug(f"Extracting codebase for files: {x['config'].get('files', [])}")
     return get_combined_docs_from_files(x["config"].get("files", []))
 
+
+# Variable definitions
+model = get_chat_model()
+llm_with_stop = model.bind(stop=["</tool_input>"])
 agent = (
         {
             "codebase": lambda x: extract_codebase(x),
@@ -108,14 +117,6 @@ agent = (
         | llm_with_stop
         | parse_output
 )
-
-
-class AgentInput(BaseModel):
-    question: str
-    config: dict = Field({})
-    chat_history: List[Tuple[str, str]] = Field(..., extra={"widget": {"type": "chat"}})
-
-
 agent_executor = AgentExecutor(
     agent=agent, tools=[], verbose=True, handle_parsing_errors=True
 ).with_types(input_type=AgentInput)
@@ -123,5 +124,4 @@ agent_executor = AgentExecutor(
 agent_executor = agent_executor | (lambda x: x["output"])
 
 if __name__ == "__main__":
-    question = "How are you ?"
-    print(agent_executor.invoke({"question": question, "chat_history": []}))
+    print(agent_executor.invoke({"question": SAMPLE_QUESTION, "chat_history": []}))

app/agents/prompts.py

@@ -183,14 +183,14 @@ Codebase ends here.
 
 Remember to strictly adhere to the Git diff format guidelines provided above when suggesting code changes.
 
+Intermediate Steps: {agent_scratchpad}
 """
-
+# Putting the intermediate steps in template as Gemini does not support
+# empty strings in API. agent_scratchpad is empty at the beginning of the chain
 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

@@ -4,10 +4,12 @@ import subprocess
 import sys
 from typing import Optional
 
+import dotenv
 from langchain_cli.cli import serve
 
-from app.utils.logging_utils import logger
 from app.utils.langchain_validation_util import validate_langchain_vars
+from app.utils.llm_constants import DEFAULT_PORT, MODEL_CHOICES
+from app.utils.logging_utils import logger
 from app.utils.version_util import get_current_version, get_latest_version
 
 
@@ -17,14 +19,16 @@ def parse_arguments():
                         help="List of files or directories to exclude (e.g., --exclude 'tst,build,*.py')")
     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", "sonnet3.5", "sonnet3.5-v2", "haiku", "opus"], default="sonnet3.5-v2",
+    parser.add_argument("--model", type=str, choices=MODEL_CHOICES, default="sonnet3.5-v2",
                         help="AWS Bedrock Model to use (e.g., --model sonnet)")
-    parser.add_argument("--port", type=int, default=6969,
+    parser.add_argument("--port", type=int, default=DEFAULT_PORT,
                         help="Port number to run Ziya frontend on (e.g., --port 8080)")
     parser.add_argument("--version", action="store_true",
                         help="Prints the version of Ziya")
     parser.add_argument("--max-depth", type=int, default=15,
                         help="Maximum depth for folder structure traversal (e.g., --max-depth 20)")
+    parser.add_argument("--env-file", type=str,
+                        help="Path to .env file containing API keys (e.g., --env-file .env)")
     return parser.parse_args()
 
 
@@ -38,7 +42,11 @@ def setup_environment(args):
         os.environ["ZIYA_AWS_PROFILE"] = args.profile
     if args.model:
         os.environ["ZIYA_AWS_MODEL"] = args.model
+        logger.info(f"Using Model: {args.model}")
     os.environ["ZIYA_MAX_DEPTH"] = str(args.max_depth)
+    if args.env_file:
+        if not dotenv.load_dotenv(args.env_file):
+            raise ValueError(f"Could not load environment file: {args.env_file}")
 
 
 def check_version_and_upgrade():

app/utils/llm_constants.py

@@ -0,0 +1,34 @@
+from typing import List, Tuple
+
+from pydantic import BaseModel, Field
+
+
+class AgentInput(BaseModel):
+    question: str
+    config: dict = Field({})
+    chat_history: List[Tuple[str, str]] = Field(..., extra={"widget": {"type": "chat"}})
+
+DEFAULT_PORT = 6969
+
+"""
+Model References: 
+- https://ai.google.dev/gemini-api/docs/models/gemini
+"""
+
+MODEL_MAPPING = {
+    "sonnet3.5": "us.anthropic.claude-3-5-sonnet-20240620-v1:0",
+    "sonnet3.5-v2": "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
+    "opus": "us.anthropic.claude-3-opus-20240229-v1:0",
+    "sonnet": "us.anthropic.claude-3-sonnet-20240229-v1:0",
+    "haiku": "us.anthropic.claude-3-haiku-20240307-v1:0",
+    "gemini-2.0-flash": "gemini-2.0-flash",
+    "gemini-2.0-pro": "gemini-2.0-pro-exp-02-05",
+    "gemini-2.0-think": "gemini-2.0-flash-thinking-exp-01-21",
+    "gemini-1.5-pro": "gemini-1.5-pro",
+    "gemini-1.5-flash": "gemini-1.5-flash",
+}
+
+MODEL_CHOICES = list(MODEL_MAPPING.keys())
+SAMPLE_QUESTION = "How are you ?"
+GEMINI_PREFIX = "gemini"
+GOOGLE_API_KEY = "GOOGLE_API_KEY"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ziya"
-version = "0.1.47"
+version = "0.1.48"
 description = ""
 authors = ["Vishnu Krishnaprasad <vishnukool@gmail.com>"]
 readme = "README.md"
@@ -15,18 +15,22 @@ uvicorn = "^0.23.2"
 pydantic = "^2.9.2"
 jinja2 = "^3.1.3"
 tiktoken = "^0.8.0"
-boto3 = "^1.34.88"
+boto3 = "1.35.75"
 langchain-aws = ">=0.2,<0.3"
 langchain = ">=0.3,<0.4"
 langgraph = ">=0.2,<0.3"
 langchainhub = ">=0.1.15"
 langchain-anthropic = ">=0.2,<0.3"
+langchain-google-genai = ">=0.1.0"
 langchain-cli = ">=0.0.15"
 pydevd-pycharm = "^243.18137.19"
 langchain-community = "^0.3.1"
+python-dotenv = "^1.0.1"
+transformers = "^4.48.3"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.3"
+black = "^23.12.1"
 
 [build-system]
 requires = ["poetry-core"]

templates/asset-manifest.json

@@ -1,7 +1,7 @@
 {
   "files": {
     "main.css": "/static/css/main.f447c982.css",
-    "main.js": "/static/js/main.59d11a1d.js",
+    "main.js": "/static/js/main.cb581e8d.js",
     "static/media/fa-solid-900.ttf": "/static/media/fa-solid-900.2582b0e4bcf85eceead0.ttf",
     "static/media/fa-brands-400.ttf": "/static/media/fa-brands-400.1815e00441357e01619e.ttf",
     "static/media/fa-solid-900.woff2": "/static/media/fa-solid-900.2463b90d9a316e4e5294.woff2",
@@ -12,10 +12,10 @@
     "static/media/fa-v4compatibility.woff2": "/static/media/fa-v4compatibility.ea8f94e1d22e0d35ccd4.woff2",
     "index.html": "/index.html",
     "main.f447c982.css.map": "/static/css/main.f447c982.css.map",
-    "main.59d11a1d.js.map": "/static/js/main.59d11a1d.js.map"
+    "main.cb581e8d.js.map": "/static/js/main.cb581e8d.js.map"
   },
   "entrypoints": [
     "static/css/main.f447c982.css",
-    "static/js/main.59d11a1d.js"
+    "static/js/main.cb581e8d.js"
   ]
 }

templates/index.html

@@ -1 +1 @@
-<!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>window.enableCodeApply="true",window.diffDisplayMode="pretty",window.diffViewType="unified"</script><script defer="defer" src="/static/js/main.59d11a1d.js"></script><link href="/static/css/main.f447c982.css" rel="stylesheet"></head><body><div id="root"></div></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>window.enableCodeApply="true",window.diffDisplayMode="pretty",window.diffViewType="unified"</script><script defer="defer" src="/static/js/main.cb581e8d.js"></script><link href="/static/css/main.f447c982.css" rel="stylesheet"></head><body><div id="root"></div></body></html>