app.py

import os

import gradio as gr
from smolagents import CodeAgent, tool

from linear_api_utils import execute_query
from sleep_per_last_token_model import SleepPerLastTokenModelLiteLLM

# .env
"""
LINEAR_API_KEY="lin_api_***"
HF_TOKEN = "hf_***"
GROQ_API_KEY = "gsk_***"
"""


def get_env_value(key, is_value_error_on_null=True):
    value = os.getenv(key)
    if value is None:
        from dotenv import load_dotenv

        load_dotenv()
        value = os.getenv(key)
        if is_value_error_on_null and value is None:
            raise ValueError(f"Need {key} on secret or .env(If running on local)")
    return value


# SETTINGS
LINEAR_ISSUE_LABEL = "huggingface-public"  # only show issue with this label,I added for demo you can remove this
## set secret key on Space setting or .env(local)
# hf_token = get_env_value("HF_TOKEN")
groq_api_key = get_env_value("GROQ_API_KEY")
api_key = get_env_value("LINEAR_API_KEY")


if api_key is None:
    raise ValueError("Need LINEAR_API_KEY on secret")
if groq_api_key is None:
    raise ValueError("Need GROQ_API_KEY on secret")


model_id = "groq/llama3-8b-8192"


def add_comment(issue_id, model_name, comment):
    comment = comment.replace('"', '\\"').replace("\n", "\\n")  # escape doublequote
    # header = f"<!---\\n start-ai-comment({model_name}) \\n--->\\n"
    header = f"[ ](start-ai-comment:{model_name})\\n"
    header += f"# {model_name.split('/')[1]}'s comment'\\n"
    comment = header + comment
    comment_create_text = """
    mutation CommentCreate {
  commentCreate(
    input: {
      issueId : "%s"
      body:"%s" 
    }
  ) {
    success
    comment {
      id
      body
    }
  }
}""" % (issue_id, comment)
    result = execute_query("add comment", comment_create_text, api_key)


issue_id = None


def change_state_reviewing():
    get_state_query_text = """
    query Sate{
    workflowStates(filter:{team:{id:{eq:"%s"}}}){
        nodes{
        id
        name
        }
    }
    }
""" % (team_id)
    result = execute_query("State", get_state_query_text, api_key)
    state_id = None
    for state in result["data"]["workflowStates"]["nodes"]:
        if state["name"] == "Reviewing":
            state_id = state["id"]
            break

    if state_id is None:
        return
    issue_update_text = """
mutation IssueUpdate {
  issueUpdate(
    id: "%s",
    input: {
      stateId: "%s",
    }
  ) {
    success
    issue {
      id
      title
      state {
        id
        name
      }
    }
  }
}
""" % (issue_id, state_id)
    result = execute_query("IssueUpdate", issue_update_text, api_key)


@tool
def get_todo_issue() -> str:
    """
    Get the Todo issue.

    Returns:
        A string describing the current issue.
    """
    global issue_id
    global issue_text
    priority_order = [1, 2, 3, 0, 4]
    for priority in priority_order:
        team_query_text = """
        query Team {
        team(id: "%s") {
            id
            issues(first:1,filter:{
                state:{
                    name:{ eq: "Todo" },
                    }
                priority:{eq:%d}
            }) {
            nodes {
                id
                title
                description
                createdAt
            }
            }
        }
        }
        
        """ % (team_id, priority)

        result = execute_query("Team", team_query_text, api_key, True)
        if len(result["data"]["team"]["issues"]["nodes"]) > 0:
            issue = result["data"]["team"]["issues"]["nodes"][0]
            issue_text = str(issue["title"])
            issue_id = issue["id"]
            description = issue.get("description", None)
            if description is not None:
                issue_text += "\n" + description
            return issue_text

    return "Not Todo issue found"


def generate_agent():
    model = SleepPerLastTokenModelLiteLLM(
        max_tokens=250,
        temperature=0.5,
        model_id=model_id,
        api_base="https://api.groq.com/openai/v1/",
        api_key=groq_api_key,
    )
    agent = CodeAgent(
        model=model,
        tools=[get_todo_issue],  ## add your tools here (don't remove final answer)
        max_steps=1,
        verbosity_level=1,
        grammar=None,
        planning_interval=None,
        name=None,
        description=None,
    )
    return agent


team_id = None


def update_text():
    def get_team_id(team_name):
        teams_text = """
    query Teams {
    teams {
        nodes {
        id
        name
        }
    }
    }
    """
        result = execute_query("Teams", teams_text, api_key)
        for team in result["data"]["teams"]["nodes"]:
            if team["name"] == team_name:
                return team["id"]
        return None

    team_name = "Agent"
    global team_id
    global issue_text
    team_id = get_team_id(team_name)

    if team_id is None:
        return f"Team {team_name} is not found", "Team not found"
    issue_text = "No Issue Found"
    agent_text = "No Agent Advice"

    agent = generate_agent()
    agent_text = agent.run(
        """
First, get the Todo using the get_todo tool.
Then, solve the Todo.
Finally, return the result of solving the Todo.
        """
    )

    add_comment(issue_id, model_id, agent_text)
    change_state_reviewing()
    #    return "", ""

    return issue_text, agent_text


with gr.Blocks() as demo:
    gr.HTML("""<h1>Linear.app API and smolagents demo</h1>
            <h2>Prepare</h2>
            <p>Need Linear.app acount and api key</a>
            <p>Remember team name and add "Reviewing" State<p> 
            """)
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Issue")
            # issue = gr.Markdown(load_text("issue.md"))
            issue = gr.Markdown("issue")
        with gr.Column():
            gr.Markdown("## Agent advice(Don't trust them completely)")
            # output = gr.Markdown(load_text("output.md"))
            output = gr.Markdown("agent result")
        demo.load(update_text, inputs=None, outputs=[issue, output])

    #
    # bt = gr.Button("Next Todo")
    # bt.click(update_text, inputs=None, outputs=[issue, output])


if __name__ == "__main__":  # without main call twice
    demo.launch()