ziya 0.1.47__tar.gz → 0.1.49__tar.gz

{ziya-0.1.47 → ziya-0.1.49}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ziya
-Version: 0.1.47
+Version: 0.1.49
 Summary: 
 Author: Vishnu Krishnaprasad
 Author-email: vishnukool@gmail.com
@@ -10,19 +10,23 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: anthropic (>=0.46.0,<0.47.0)
 Requires-Dist: boto3 (>=1.34.88,<2.0.0)
 Requires-Dist: jinja2 (>=3.1.3,<4.0.0)
 Requires-Dist: langchain (>=0.3,<0.4)
 Requires-Dist: langchain-anthropic (>=0.2,<0.3)
 Requires-Dist: langchain-aws (>=0.2,<0.3)
-Requires-Dist: langchain-cli (>=0.0.15)
+Requires-Dist: langchain-cli (>=0.0.35,<0.0.36)
 Requires-Dist: langchain-community (>=0.3.1,<0.4.0)
 Requires-Dist: langchainhub (>=0.1.15)
 Requires-Dist: langgraph (>=0.2,<0.3)
+Requires-Dist: patch-ng (>=1.17)
 Requires-Dist: pydantic (>=2.9.2,<3.0.0)
 Requires-Dist: pydevd-pycharm (>=243.18137.19,<244.0.0)
 Requires-Dist: tiktoken (>=0.8.0,<0.9.0)
+Requires-Dist: transformers (>=4.36.2,<5.0.0)
 Requires-Dist: uvicorn (>=0.23.2,<0.24.0)
+Requires-Dist: whatthepatch (>=1.0.2)
 Description-Content-Type: text/markdown
 
 # Ziya
@@ -84,13 +88,13 @@ ziya
 
 `--profile`: AWS profile to use for the Bedrock LLM.
 
-`--model`: The AWS Bedrock Model to use, one of `sonnet3.5`(default), `sonnet`, `haiku` or `opus`.
+`--model`: The AWS Bedrock Model to use, one of `sonnet3.7`(default), `sonnet`, `haiku` or `opus`.
 
 `--port`: The port number for frontend app. Default is `6969`.
 
 `--max-depth`: Maximum depth for folder structure traversal. Default is `15`.
 
 ```bash
-ziya --exclude='tst,build,*.py' --profile=ziya --model=sonnet3.5 --port=8080
+ziya --exclude='tst,build,*.py' --profile=ziya --model=sonnet3.7 --port=8080
 ```
 

{ziya-0.1.47 → ziya-0.1.49}/README.md

@@ -57,12 +57,12 @@ ziya
 
 `--profile`: AWS profile to use for the Bedrock LLM.
 
-`--model`: The AWS Bedrock Model to use, one of `sonnet3.5`(default), `sonnet`, `haiku` or `opus`.
+`--model`: The AWS Bedrock Model to use, one of `sonnet3.7`(default), `sonnet`, `haiku` or `opus`.
 
 `--port`: The port number for frontend app. Default is `6969`.
 
 `--max-depth`: Maximum depth for folder structure traversal. Default is `15`.
 
 ```bash
-ziya --exclude='tst,build,*.py' --profile=ziya --model=sonnet3.5 --port=8080
+ziya --exclude='tst,build,*.py' --profile=ziya --model=sonnet3.7 --port=8080
 ```

ziya-0.1.49/app/agents/agent.py

@@ -0,0 +1,389 @@
+import os
+import os.path
+from typing import Dict, List, Tuple, Set, Union, Optional, Any
+
+import json
+import time
+import botocore
+from langchain.agents import AgentExecutor
+from langchain.agents.format_scratchpad import format_xml
+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 langchain_core.outputs import Generation
+from langchain_core.runnables import RunnablePassthrough, Runnable
+from pydantic import BaseModel, Field
+
+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.file_utils import is_binary_file
+from app.utils.file_state_manager import FileStateManager
+
+import tiktoken
+import anthropic
+from anthropic import Anthropic
+
+
+def clean_chat_history(chat_history: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+    """Clean chat history by removing invalid messages and normalizing content."""
+    cleaned = []
+    for human, ai in chat_history:
+        # Skip pairs with empty messages
+        if not human or not human.strip() or not ai or not ai.strip():
+            logger.warning(f"Skipping invalid message pair: human='{human}', ai='{ai}'")
+            continue
+        cleaned.append((human.strip(), ai.strip()))
+    return cleaned
+
+def _format_chat_history(chat_history: List[Tuple[str, str]]) -> List[Union[HumanMessage, AIMessage]]:
+    logger.info(f"Formatting chat history: {json.dumps(chat_history, indent=2)}")
+    cleaned_history = clean_chat_history(chat_history)
+    buffer = []
+    for human, ai in cleaned_history:
+        buffer.append(HumanMessage(content=human))
+        buffer.append(AIMessage(content=ai))
+    return buffer
+
+def parse_output(message):
+    """Parse and sanitize the output from the language model."""
+    try:
+        # Extract content from message object
+        if isinstance(message, str):
+            content = message
+        else:
+            content = getattr(message, 'text', None) or getattr(message, 'content', None) or str(message)
+
+        if isinstance(content, str):
+            # Check if this is an error message
+            try:
+                error_data = json.loads(content)
+                if error_data.get('error') == 'validation_error':
+                    logger.info(f"Detected validation error in output: {content}")
+                    return AgentFinish(return_values={"output": content}, log=content)
+            except (json.JSONDecodeError, TypeError):
+                pass
+        # If not an error, clean and return the content
+        text = clean_backtick_sequences(str(content))
+        logger.info(f"parse_output received content size: {len(str(content))} chars, returning size: {len(text)} chars")
+        return AgentFinish(return_values={"output": text}, log=text)
+    except Exception as e:
+        logger.error(f"Error in parse_output: {str(e)}")
+        return AgentFinish(return_values={"output": str(message)}, log=str(message))
+
+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.7": "us.anthropic.claude-3-7-sonnet-20250219-v1:0",
+    "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, 
+        retries={ 
+            'max_attempts': 3, 
+            'total_max_attempts': 5 
+        } 
+        # retry_mode is not supported in this version
+    )
+)
+
+# Create a wrapper class that adds retries
+class RetryingChatBedrock(Runnable):
+    def __init__(self, model):
+        self.model = model
+
+    def bind(self, **kwargs):
+        return RetryingChatBedrock(self.model.bind(**kwargs))
+
+    def get_num_tokens(self, text: str) -> int:
+        """
+        Custom token counting function using anthropic's tokenizer.
+        Falls back to tiktoken if anthropic is not available.
+        """
+        try:
+            client = Anthropic()
+            count = client.count_tokens(text)
+            return count
+        except Exception as e:
+            logger.warning(f"Failed to use anthropic tokenizer: {str(e)}")
+            try:
+                return len(tiktoken.get_encoding("cl100k_base").encode(text))
+            except Exception as e:
+                logger.error(f"Failed to count tokens: {str(e)}")
+                return len(text.split()) # Rough approximation
+
+    def __getattr__(self, name: str):
+        # Delegate any unknown attributes to the underlying model
+        return getattr(self.model, name)
+
+    async def _handle_stream_error(self, e: Exception):
+        """Handle stream errors by yielding an error message."""
+        yield Generation(
+            text=json.dumps({
+                "error": "validation_error",
+                "detail": "Selected content is too large for the model. Please reduce the number of files."
+            })
+        )
+        return
+
+    async def _handle_validation_error(self, e: Exception):
+        """Handle validation errors by yielding an error message."""
+        error_chunk = Generation(
+            text=json.dumps({"error": "validation_error", "detail": str(e)})
+        )
+        yield error_chunk
+    def _is_streaming(self, func) -> bool:
+        """Check if this is a streaming operation."""
+        return hasattr(func, '__name__') and func.__name__ == 'astream'
+
+    def invoke(self, input: Any, config: Optional[Dict] = None, **kwargs) -> Any:
+        return self.model.invoke(input, config, **kwargs)
+
+    async def ainvoke(self, input: Any, config: Optional[Dict] = None, **kwargs) -> Any:
+        return await self.model.ainvoke(input, config, **kwargs)
+
+    async def astream(self, input: Any, config: Optional[Dict] = None, stream_mode: bool = True, **kwargs):
+        max_retries = 3
+        retry_delay = 1
+
+        for attempt in range(max_retries):
+            try:
+                if stream_mode:
+                    async for chunk in self.model.astream(input, config, **kwargs):
+                        if isinstance(chunk, dict) and "error" in chunk:
+                            yield Generation(text=json.dumps(chunk))
+                        else:
+                            yield chunk
+                else:
+                    # Get complete response at once
+                    result = await self.model.ainvoke(input, config, **kwargs)
+                    yield result
+                break
+            except (botocore.exceptions.EventStreamError, ExceptionGroup) as e:
+                error_message = str(e)
+                if "validationException" in str(e):
+                    logger.error(f"Validation error from AWS: {str(e)}")
+                    yield Generation(text=json.dumps({
+                        "error": "validation_error",
+                        "detail": "Selected content is too large for the model. Please reduce the number of files."
+                    }))
+                    return
+            raise e
+
+
+model = RetryingChatBedrock(model)
+
+file_state_manager = FileStateManager()
+
+def get_combined_docs_from_files(files, conversation_id: str = "default") -> str:
+    combined_contents: str = ""
+    logger.debug("Processing files:")
+    print_file_tree(files if isinstance(files, list) else files.get("config", {}).get("files", []))
+    logger.info(f"Processing files with conversation_id: {conversation_id}")
+
+    user_codebase_dir: str = os.environ["ZIYA_USER_CODEBASE_DIR"]
+    for file_path in files:
+        full_path = os.path.join(user_codebase_dir, file_path)
+        # Skip directories
+        if os.path.isdir(full_path):
+            continue
+        try:
+            # Get annotated content with change tracking
+            annotated_lines, success = file_state_manager.get_annotated_content(conversation_id, file_path)
+            if success:
+                combined_contents += f"File: {file_path}\n" + "\n".join(annotated_lines) + "\n\n"
+        except Exception as e:
+            logger.error(f"Error processing {file_path}: {str(e)}")
+
+    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
+
+
+llm_with_stop = model.bind(stop=["</tool_input>"])
+
+class AgentInput(BaseModel):
+    question: str
+    config: dict = Field({})
+    chat_history: List[Tuple[str, str]] = Field(..., extra={"widget": {"type": "chat"}})
+    conversation_id: str = Field(default="default", description="Unique identifier for the conversation")
+
+def extract_codebase(x):
+    files = x["config"].get("files", [])
+    conversation_id = x.get("conversation_id", "default")
+    logger.debug(f"Extracting codebase for files: {files}")
+    logger.info(f"Processing with conversation_id: {conversation_id}")
+
+    file_contents = {}
+    for file_path in files:
+
+        try:
+            full_path = os.path.join(os.environ["ZIYA_USER_CODEBASE_DIR"], file_path)
+            if os.path.isdir(full_path):
+                logger.debug(f"Skipping directory: {file_path}")
+                continue
+            if is_binary_file(full_path):
+                logger.debug(f"Skipping binary file: {file_path}")
+                continue
+
+            with open(full_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+                file_contents[file_path] = content
+                logger.info(f"Successfully loaded {file_path} with {len(content.splitlines())} lines")
+        except (UnicodeDecodeError, IOError) as e:
+                logger.error(f"Error reading file {file_path}: {str(e)}")
+                continue
+
+    # Initialize or update file states
+    if conversation_id not in file_state_manager.conversation_states:
+        file_state_manager.initialize_conversation(conversation_id, file_contents)
+
+    # Update any new files that weren't in the initial state
+    file_state_manager.update_files_in_state(conversation_id, file_contents)
+
+    # Get changes since last message
+    overall_changes, recent_changes = file_state_manager.format_context_message(conversation_id)
+
+    codebase = get_combined_docs_from_files(files, conversation_id)
+    logger.info(f"Changes detected - Overall: {bool(overall_changes)}, Recent: {bool(recent_changes)}")
+
+    result = []
+
+    # Add recent changes first if any
+    if recent_changes:
+        result.append(recent_changes)
+        result.append("")
+
+    # Add overall changes if any
+    if overall_changes:
+        result.extend([
+            "SYSTEM: Overall Code Changes",
+            "------------------------",
+            overall_changes
+        ])
+        result.append("")
+
+    # Add the codebase content
+    result.append(codebase)
+
+    final_string = "\n".join(result)
+    file_markers = [line for line in final_string.split('\n') if line.startswith('File: ')]
+    logger.info(f"Final string assembly:")
+    logger.info(f"Total length: {len(final_string)} chars")
+    logger.info(f"Number of File: markers: {len(file_markers)}")
+    logger.info(f"First 500 chars:\n{final_string[:500]}")
+    logger.info(f"Last 500 chars:\n{final_string[-500:]}")
+
+    # Debug the content at each stage
+    logger.info("Content flow tracking:")
+    logger.info(f"1. Number of files in file_contents: {len(file_contents)}")
+    logger.info(f"2. Number of files in conversation state: {len(file_state_manager.conversation_states.get(conversation_id, {}))}")
+
+    # Check content before joining
+    file_headers = [line for line in codebase.split('\n') if line.startswith('File: ')]
+    logger.info(f"3. Files in codebase string:\n{chr(10).join(file_headers)}")
+
+    logger.info(f"Final assembled context length: {len(result)} sections, {sum(len(s) for s in result)} total characters")
+    file_headers = [line for line in codebase.split('\n') if line.startswith('File: ')]
+    logger.info(f"Number of files in codebase: {len(file_headers)}")
+    if file_headers:
+        logger.info(f"First few files in codebase:\n{chr(10).join(file_headers[:5])}")
+
+    if result:
+        return final_string
+    return codebase
+
+def log_output(x):
+    logger.info(f"Final output size: {len(x.return_values['output'])} chars, first 100 chars: {x.return_values['output'][:100]}")
+    return x
+
+def log_codebase_wrapper(x):
+    codebase = extract_codebase(x)
+    logger.info(f"Codebase before prompt: {len(codebase)} chars")
+    file_count = len([l for l in codebase.split('\n') if l.startswith('File: ')])
+    logger.info(f"Number of files in codebase before prompt: {file_count}")
+    logger.info(f"Files in codebase before prompt:{chr(10)}{chr(10).join([l for l in codebase.split(chr(10)) if l.startswith('File: ')])}")
+    return codebase
+
+# Define the agent chain
+agent = (
+    {
+        "codebase": log_codebase_wrapper,
+        "question": lambda x: x["question"],
+        "agent_scratchpad": lambda x: format_xml(x["intermediate_steps"]),
+        "chat_history": lambda x: _format_chat_history(x["chat_history"]),
+    }
+    | conversational_prompt
+    | (lambda x: (
+        logger.info(f"Template population check:") or
+        logger.info(f"System message contains codebase section: {'---------------------------------------' in str(x)}") or
+        logger.info(f"Number of 'File:' markers in system message: {str(x).count('File:')}") or
+        x))
+    | llm_with_stop
+    | parse_output
+    | log_output
+)
+
+def update_conversation_state(conversation_id: str, file_paths: List[str]) -> None:
+    """Update file states after a response has been generated"""
+    logger.info(f"Updating conversation state for {conversation_id} with {len(file_paths)} files")
+    # Read current file contents, skipping directories
+    file_contents = {}
+    for file_path in file_paths:
+        full_path = os.path.join(os.environ["ZIYA_USER_CODEBASE_DIR"], file_path)
+        if not os.path.isdir(full_path):
+            try:
+                with open(full_path, 'r', encoding='utf-8') as f:
+                    file_contents[file_path] = f.read()
+                logger.debug(f"Read current content for {file_path}")
+            except Exception as e:
+                logger.error(f"Error reading file {file_path}: {str(e)}")
+                continue
+
+    # Update states and get changes
+    changes = file_state_manager.update_files(conversation_id, file_contents)
+    logger.info(f"File state update complete. Changes detected: {bool(changes)}")
+    if changes:
+        logger.info("Changes detected during update:")
+        logger.info(json.dumps(changes, indent=2))
+        logger.info(f"Files changed during conversation {conversation_id}:")
+        for file_path, changed_lines in changes.items():
+            logger.info(f"- {file_path}: {len(changed_lines)} lines changed")
+
+def update_and_return(input_data: Dict[str, Any]) -> Dict[str, Any]:
+    """Update file state and preserve the full response structure"""
+    update_conversation_state(input_data.get("conversation_id", "default"),
+                            input_data.get("config", {}).get("files", []))
+    return input_data
+
+# Finally create the executor
+agent_executor = AgentExecutor(
+    agent=agent,
+    tools=[],
+    verbose=True,
+    handle_parsing_errors=True,
+    max_iterations=3
+).with_types(input_type=AgentInput) | RunnablePassthrough(update_and_return)
+
+# Chain the executor with the state update
+agent_executor = agent_executor | RunnablePassthrough(update_and_return)

{ziya-0.1.47 → ziya-0.1.49}/app/agents/prompts.py

@@ -15,6 +15,28 @@ CRITICAL: INSTRUCTION PRESERVATION:
 
 You are an excellent coder. Help the user with their coding tasks. You are given the codebase of the user in your context.
 
+IMPORTANT: Code Context Format
+
+The codebase context includes line-by-line change indicators in this format:
+[NNN+] - Line number NNN is newly added since conversation began
+[NNN*] - Line number NNN was modified since conversation began
+[NNN ] - Line number NNN is unchanged (note the space)
+
+Example:
+File: example.py
+[001+] def new_function():      # This is a newly added line
+[002*]     x = 42              # This line was modified
+[003 ]     return x           # This line is unchanged
+
+When you see a "Code Change Summary" at the start of the context, it indicates files
+that have been modified during our conversation. Use this information to maintain
+context about the evolution of the code during our discussion.
+
+CRITICAL: Every marked change (-/+) must show actual content differences:
+ Never output identical content as a change, even if spacing differs
+ Skip single-line changes that differ only in non-functional whitespace
+ Only include whitespace changes that affect functionality (e.g., Python indentation)
+
 IMPORTANT: When recommending code changes, format your response as a standard Git diff format unless the user specifies otherwise. 
 Follow these strict guidelines for diff formatting:
 
@@ -34,20 +56,13 @@ Follow these strict guidelines for diff formatting:
    - Use + for each line of the new file content.
       - Always count lines in the original file before generating the diff
    - Include context that identifies the location unambiguously:
-     * For CSS: Always include the selector in context
-       Good:
-       ```diff
-       @@ -130,6 +130,7 @@ .folder-tree-panel
-         position: fixed;
-       ```
-       Bad:
-       ```diff
-       @@ -131,6 +131,7 @@
-         position: fixed;
-       ```
+     * For CSS: Always include the complete selector line WITH its opening brace as a context line 
+       The selector line with its opening brace MUST be included even if it means
+           adding extra context lines
      * For functions: Include the function declaration
      * For classes: Include the class declaration
      * For nested blocks: Include parent identifier
+     * Merge chunks separated by fewer than 5 unchanged lines into a single chunk
   - When counting lines for @@ markers:
      * First number pair (-A,B) refers to the original file
      * Second number pair (+C,D) refers to the new file
@@ -70,18 +85,29 @@ CRITICAL: When generating hunks and context:
    - Must include the identifying name/selector/declaration
    - Don't need the entire block, just enough for identification
    - For nested items, include immediate parent identifier
-   - Prefer starting at a named block boundary when possible
+   - Always include sufficient context to locate the change:
+    * When modifying within a named block (function, class, configuration, etc.):
+        - Show the block's declaration/name and opening delimiter
+        - Or show the closing delimiter and the next uniquely named block's opening line
+      * When modifying near block boundaries:
+        - Include both the ending of the previous block and start of the next
+        - Ensure each block shown has its identifying name/declaration visible
+      * For nested structures:
+        - Show enough parent context to uniquely identify the location
+        - Include at least one uniquely identifying parent name/declaration
+      * Always ensure the context makes the location unambiguously identifiable
 3. Verify line numbers match the actual file content
 4. Double-check that context lines exist in the original file
 
 CRITICAL: VISUALIZATION CAPABILITIES:
-You can generate inline diagrams using ```graphviz code blocks. Actively look for
-opportunities to enhance explanations with visual representations when they would
-provide clearer understanding, especially for:
+You can generate inline diagrams using either ```graphviz or ```d3 code blocks. 
+Actively look for opportunities to enhance explanations with visual representations 
+when they would provide clearer understanding, especially for:
 - System architectures
 - Flow diagrams
 - Dependency relationships
 - Complex structures or processes
+Use graphviz for architecture and flow diagrams
 
 IMPORTANT: When making changes:
 1. Focus only on fixing the specific problem described by the user
@@ -188,7 +214,7 @@ Remember to strictly adhere to the Git diff format guidelines provided above whe
 conversational_prompt = ChatPromptTemplate.from_messages(
     [
         ("system", template),
-        # ("system", "You are a helpful AI bot. Your name is {name}."),
+        # ("system", "You are a helpful AI bot."),
         MessagesPlaceholder(variable_name="chat_history"),
         ("user", "{question}"),
         ("ai", "{agent_scratchpad}"),

{ziya-0.1.47 → ziya-0.1.49}/app/main.py

@@ -17,7 +17,7 @@ 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=["sonnet", "sonnet3.5", "sonnet3.7", "sonnet3.5-v2", "haiku", "opus"], default="sonnet3.5-v2",
                         help="AWS Bedrock Model to use (e.g., --model sonnet)")
     parser.add_argument("--port", type=int, default=6969,
                         help="Port number to run Ziya frontend on (e.g., --port 8080)")