Upload folder using huggingface_hub

app.py

@@ -4,11 +4,14 @@ import asyncio
 import logging
 import json
 import requests
+import uuid
 from datetime import datetime
+from typing import Literal, TypedDict, List, Optional, Tuple, Any
+from pydantic import BaseModel, Field
 from langchain_openai import ChatOpenAI
 from langchain.agents import AgentExecutor, create_openai_functions_agent
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
 from langchain_mcp_adapters.tools import load_mcp_tools
 from mcp import StdioServerParameters
 from tools_api import search_web_tool, search_hf_spaces_tool
@@ -16,6 +19,13 @@ from langchain.tools import Tool
 from git import Repo
 import shutil
 
+# LangGraph imports
+from langgraph.graph import StateGraph, START, END, MessagesState
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.store.memory import InMemoryStore
+from langgraph.store.base import BaseStore
+from langgraph.types import Command
+
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("agent-app")
@@ -23,9 +33,9 @@ logger = logging.getLogger("agent-app")
 # Constants
 JULES_FILES_REPO = "https://github.com/JsonLord/jules_files.git"
 WORK_DIR = "/tmp/agent_work"
+HELMHOLTZ_BASE_URL = "https://api.helmholtz-blablador.fz-juelich.de/v1"
 
 # Initialize LLMs (Helmholtz Blablador)
-HELMHOLTZ_BASE_URL = "https://api.helmholtz-blablador.fz-juelich.de/v1"
 api_key = os.environ.get("BLABLADOR_API_KEY")
 
 chat_llm = ChatOpenAI(
@@ -49,6 +59,70 @@ fast_llm = ChatOpenAI(
     max_tokens=512
 )
 
+# Memory Stores
+checkpointer = MemorySaver()
+long_term_store = InMemoryStore()
+
+# Schemas for Memory
+class UserPreferences(BaseModel):
+    """Updated user preferences and project context."""
+    chain_of_thought: str = Field(description="Reasoning about what needs to be remembered.")
+    user_preferences: str = Field(description="Updated user preferences and context.")
+
+class State(MessagesState):
+    """Central state for the agent graph."""
+    project_state: str # JSON string of project_state
+
+# Memory Functions
+def get_memory(store, namespace, key="user_preferences", default_content=""):
+    """Retrieve memory from store or initialize with default."""
+    memory = store.get(namespace, key)
+    if memory:
+        return memory.value
+    else:
+        store.put(namespace, key, default_content)
+        return default_content
+
+MEMORY_UPDATE_INSTRUCTIONS = """
+# Role
+You are a memory profile manager for an AI application assistant.
+
+# Rules
+- NEVER overwrite the entire profile
+- ONLY add new information or update facts contradicted by feedback
+- PRESERVE all other information
+- Focus on user preferences, project goals, and technical decisions.
+
+# Process current profile for {namespace}
+<memory_profile>
+{current_profile}
+</memory_profile>
+"""
+
+async def update_memory(store, namespace, messages, current_profile):
+    """Intelligently update the memory store based on conversation."""
+    if not messages:
+        return current_profile
+        
+    memory_updater_llm = fast_llm.with_structured_output(UserPreferences)
+    
+    messages_to_send = [
+        msg for msg in messages
+        if isinstance(msg, (HumanMessage, AIMessage))
+    ]
+    
+    try:
+        result = await memory_updater_llm.ainvoke(
+            [
+                {"role": "system", "content": MEMORY_UPDATE_INSTRUCTIONS.format(current_profile=current_profile, namespace=namespace)},
+            ] + messages_to_send
+        )
+        store.put(namespace, "user_preferences", result.user_preferences)
+        return result.user_preferences
+    except Exception as e:
+        logger.error(f"Failed to update memory: {e}")
+        return current_profile
+
 # Shared memory and context management
 def convert_to_langchain_messages(history):
     """Convert Gradio history (list of dicts) to LangChain message objects."""
@@ -61,24 +135,13 @@ def convert_to_langchain_messages(history):
                 messages.append(HumanMessage(content=content))
             elif role == "assistant":
                 messages.append(AIMessage(content=content))
-        elif isinstance(msg, (list, tuple)) and len(msg) == 2:
-            # Handle legacy tuple format just in case
-            messages.append(HumanMessage(content=msg[0]))
-            messages.append(AIMessage(content=msg[1]))
     return messages
 
-def truncate_history(history, max_messages=5):
-    """Truncate chat history to fit smaller context windows."""
-    if not history:
-        return []
-    return history[-max_messages:]
-
 # Log required secrets reminder
 required_secrets = [
     "BLABLADOR_API_KEY", 
     "GITHUB_TOKEN", 
     "JULES_API_KEY", 
-    "CONTEXTSTREAM_API_KEY",
     "HF_TOKEN",
     "HF_PROFILE"
 ]
@@ -90,15 +153,14 @@ for secret in required_secrets:
         logger.warning(f"[MISSING] {secret} is not set. Please add it to your Space secrets.")
 logger.info("==========================================")
 
-# Memory/Storage
-project_state = {
+# Initial Project State
+initial_project_state = {
     "idea": "",
     "overview": "",
     "repo_url": ""
 }
 
 async def get_all_tools():
-    # 1. Custom Search Tools
     custom_tools = [
         Tool(
             name="search_web",
@@ -112,75 +174,87 @@ async def get_all_tools():
         )
     ]
     
-    # 2. MCP Tools
     mcp_tools = []
     
-    # Jules MCP
-    try:
-        jules_mcp = await load_mcp_tools(StdioServerParameters(
-            command="python3",
-            args=["mcp/mcp_jules.py"],
-            env=os.environ.copy()
-        ))
-        mcp_tools.extend(jules_mcp)
-    except Exception as e:
-        logger.error(f"Failed to load Jules MCP tools: {e}")
-
-    # ContextStream MCP
+    # Try loading MCP tools
     try:
-        cs_mcp = await load_mcp_tools(StdioServerParameters(
-            command="npx",
-            args=["-y", "@contextstream/mcp-server"],
-            env=os.environ.copy()
-        ))
-        mcp_tools.extend(cs_mcp)
+        # Jules
+        j_params = StdioServerParameters(command="python3", args=["mcp/mcp_jules.py"], env=os.environ.copy())
+        # The correct way might be passing a list of servers to a manager
+        # For now, let's stick to what we have or just skip them if they fail
+        mcp_tools.extend(await load_mcp_tools(j_params))
     except Exception as e:
-        logger.error(f"Failed to load ContextStream MCP tools: {e}")
+        logger.warning(f"Jules MCP skipped: {e}")
 
-    # GitHub MCP
     try:
-        github_mcp = await load_mcp_tools(StdioServerParameters(
-            command="npx",
-            args=["-y", "@modelcontextprotocol/server-github"],
-            env=os.environ.copy()
-        ))
-        mcp_tools.extend(github_mcp)
+        # GitHub
+        gh_params = StdioServerParameters(command="npx", args=["-y", "@modelcontextprotocol/server-github"], env=os.environ.copy())
+        mcp_tools.extend(await load_mcp_tools(gh_params))
     except Exception as e:
-        logger.error(f"Failed to load GitHub MCP tools: {e}")
+        logger.warning(f"GitHub MCP skipped: {e}")
 
     return custom_tools + mcp_tools
 
-async def create_agent(llm_instance):
+# LangGraph Nodes
+async def agent_node(state: State, config: Any, store: BaseStore):
+    """Agent reasoning node."""
     tools = await get_all_tools()
+    
+    # Retrieve LTM
+    user_id = config.get("configurable", {}).get("user_id", "default_user")
+    namespace = (user_id, "ideate_memory")
+    ltm_context = get_memory(store, namespace, default_content="No previous context.")
+    
     prompt = ChatPromptTemplate.from_messages([
-        ("system", "You are an autonomous agent assisting in building and improving applications. "
-                   "Use your tools to search the web, explore HF spaces, manage GitHub repos, and interact with Jules. "
-                   "Current project state: {project_state}"),
-        MessagesPlaceholder(variable_name="chat_history"),
-        ("human", "{input}"),
+        ("system", "You are an autonomous agent assisting in building applications. "
+                   "Use tools to search web/HF, manage GitHub, and Jules. "
+                   "\nLong-Term Memory/Context: {ltm_context}"
+                   "\nProject State: {project_state}"),
+        MessagesPlaceholder(variable_name="messages"),
         MessagesPlaceholder(variable_name="agent_scratchpad"),
     ])
-    agent = create_openai_functions_agent(llm_instance, tools, prompt)
-    return AgentExecutor(agent=agent, tools=tools, verbose=True)
+    
+    agent = create_openai_functions_agent(chat_llm, tools, prompt)
+    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+    
+    try:
+        result = await agent_executor.ainvoke({
+            "messages": state["messages"],
+            "ltm_context": ltm_context,
+            "project_state": state["project_state"]
+        })
+        output = result["output"]
+    except Exception as e:
+        logger.error(f"Agent execution failed: {e}")
+        output = f"I encountered an error: {e}"
+    
+    # Update memory after each turn
+    await update_memory(store, namespace, state["messages"] + [AIMessage(content=output)], ltm_context)
+    
+    return {"messages": [AIMessage(content=output)]}
+
+# Assemble Graph
+workflow = StateGraph(State)
+workflow.add_node("agent", agent_node)
+workflow.add_edge(START, "agent")
+workflow.add_edge("agent", END)
+compiled_graph = workflow.compile(checkpointer=checkpointer, store=long_term_store)
 
 # Session handlers
 async def handle_ideate(message, history):
-    # Use chat_llm for ideation
-    agent_executor = await create_agent(chat_llm)
+    config = {"configurable": {"thread_id": "global_thread", "user_id": "user_1"}}
+    lc_messages = convert_to_langchain_messages(history)
+    lc_messages.append(HumanMessage(content=message))
     
-    # Convert Gradio history to LangChain messages
-    lc_history = convert_to_langchain_messages(history)
+    p_state = json.dumps(initial_project_state)
     
-    # Large model can handle more context, but still good to be mindful
-    processed_history = truncate_history(lc_history, max_messages=10)
+    result = await compiled_graph.ainvoke(
+        {"messages": lc_messages, "project_state": p_state},
+        config=config
+    )
     
-    result = await agent_executor.ainvoke({
-        "input": message,
-        "chat_history": processed_history,
-        "project_state": json.dumps(project_state)
-    })
-    project_state["idea"] += f"\nUser: {message}\nAgent: {result['output']}"
-    return result["output"]
+    bot_message = result["messages"][-1].content
+    return bot_message
 
 async def handle_github_prep(idea_description, target_repo_name):
     if not idea_description:
@@ -189,17 +263,16 @@ async def handle_github_prep(idea_description, target_repo_name):
         return "Please provide a target repository name."
     
     try:
-        # 1. Clone jules_files
         repo_path = os.path.join(WORK_DIR, "jules_files")
         if os.path.exists(repo_path):
             shutil.rmtree(repo_path)
         Repo.clone_from(JULES_FILES_REPO, repo_path)
         
-        # 2. Adapt files
         to_be_adapted_dir = os.path.join(repo_path, "to_be_adapted")
         jules_temp_dir = os.path.join(repo_path, "jules")
         os.makedirs(jules_temp_dir, exist_ok=True)
         
+        # Adapt existing files
         for filename in os.listdir(to_be_adapted_dir):
             file_path = os.path.join(to_be_adapted_dir, filename)
             if os.path.isfile(file_path):
@@ -208,8 +281,21 @@ async def handle_github_prep(idea_description, target_repo_name):
                 response = await code_llm.ainvoke(f"Adapt the following file according to instructions for the project idea: {idea_description}\n\nInstructions:\n{instructions}")
                 with open(os.path.join(jules_temp_dir, filename), 'w') as f:
                     f.write(response.content)
+
+        # Explicitly generate AGENTS.md for Jules behavior and context
+        agents_md_prompt = f"""Generate an AGENTS.md file for the project: {idea_description}. 
+        This file is for 'Google Jules' (an AI coding agent).
+        It must contain:
+        1. Working behavior: How Jules should approach tasks in this repo.
+        2. Prompt context: Key information Jules must keep in context.
+        3. Instructions on how to follow the other project files (Project_Overview, etc.).
+        4. Tips for Jules to achieve the best results for this specific project.
+        Format it in Markdown."""
+        
+        agents_response = await code_llm.ainvoke(agents_md_prompt)
+        with open(os.path.join(jules_temp_dir, "AGENTS.md"), 'w') as f:
+            f.write(agents_response.content)
         
-        # 3. Push to target repo
         github_token = os.environ.get("GITHUB_TOKEN")
         if not github_token:
             return "GITHUB_TOKEN not found in environment."
@@ -242,18 +328,28 @@ async def handle_github_prep(idea_description, target_repo_name):
         repo.index.commit(f"Add adapted jules files for project: {idea_description[:50]}...")
         repo.git.push()
         
-        project_state["repo_url"] = f"https://github.com/JsonLord/{target_repo_name}"
         return f"Files adapted and pushed to JsonLord/{target_repo_name}/jules."
     except Exception as e:
         logger.error(f"GitHub Prep failed: {e}")
         return f"Error: {str(e)}"
 
 async def handle_jules_comm(repo_url):
-    agent_executor = await create_agent(fast_llm)
+    agent_executor = await create_agent_executor(fast_llm)
     prompt = f"Create a new session in Jules for the repo {repo_url} and start implementing based on the files in /jules folder."
     result = await agent_executor.ainvoke({"input": prompt, "chat_history": [], "project_state": ""})
     return result["output"]
 
+async def create_agent_executor(llm_instance):
+    tools = await get_all_tools()
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", "You are an assistant for Jules session management."),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{input}"),
+        MessagesPlaceholder(variable_name="agent_scratchpad"),
+    ])
+    agent = create_openai_functions_agent(llm_instance, tools, prompt)
+    return AgentExecutor(agent=agent, tools=tools, verbose=True)
+
 async def handle_test():
     results = []
     try:

requirements.txt

@@ -1,7 +1,8 @@
 gradio
-langchain==0.3.17
-langchain-openai==0.2.14
-langchain-mcp-adapters==0.1.0
+langchain
+langchain-openai
+langchain-mcp-adapters
+langgraph
 requests
 python-dotenv
 huggingface_hub