Aryan Khurana's Blog
Implementing a feature in someone else's project

Implementing a feature in someone else's project

Enhancing a CLI tool with token usage tracking while navigating the open source contribution process

Aryan Khurana's photo
Aryan Khurana
·

6 min read

Table of contents

Introduction

In our open source course at Seneca this week, we were tasked with learning the full open source contribution flow by collaborating on someone else’s repository and adding a feature. I chose to contribute to my friend Harshil’s project, ResumeEnhancer. This CLI tool is designed to optimize and tailor resumes based on specific job descriptions.

The feature I am working on is Token Usage Reporting. This feature will allow users to track the number of tokens sent and received in each interaction with the LLM (Language Model). By adding a command-line flag (--token-usage or tu), users will be able to receive detailed token usage statistics for their resume enhancement requests.

The Workflow

Creating an Issue

The first step in the contribution process was to create an issue in Harshil’s repository outlining the new feature. You can view the issue I created here. The issue details what is expected for this feature and serves as the reference for my implementation.

Understanding the codebase

After creating the issue, I focused on understanding how Harshil’s code is structured and formatted. This involved:

  • Running the code locally to see how it operates.

  • Reviewing the file structure to understand the role of each file.

  • Investigating how the CLI tool processes resumes and integrates with the groq API.

Learning About groq API

Since the project utilizes the groq API, I spent some time learning how to interact with this API. This knowledge is crucial for implementing the token usage feature effectively. I reviewed the API reference to understand how to fetch token usage data and integrate it into the existing codebase.

This is what a response from grok looks like so my task was to simply extract the usage information.

{
  "id": "34a9110d-c39d-423b-9ab9-9c748747b204",
  "object": "chat.completion",
  "created": 1708045122,
  "model": "mixtral-8x7b-32768",
  "system_fingerprint": null,
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": "Low latency Large Language Models (LLMs) are important in the field of artificial intelligence and natural language processing (NLP) for several reasons:\n\n1. Real-time applications: Low latency LLMs are essential for real-time applications such as chatbots, voice assistants, and real-time translation services. These applications require immediate responses, and high latency can lead to a poor user experience.\n\n2. Improved user experience: Low latency LLMs provide a more seamless and responsive user experience. Users are more likely to continue using a service that provides quick and accurate responses, leading to higher user engagement and satisfaction.\n\n3. Competitive advantage: In today's fast-paced digital world, businesses that can provide quick and accurate responses to customer inquiries have a competitive advantage. Low latency LLMs can help businesses respond to customer inquiries more quickly, potentially leading to increased sales and customer loyalty.\n\n4. Better decision-making: Low latency LLMs can provide real-time insights and recommendations, enabling businesses to make better decisions more quickly. This can be particularly important in industries such as finance, healthcare, and logistics, where quick decision-making can have a significant impact on business outcomes.\n\n5. Scalability: Low latency LLMs can handle a higher volume of requests, making them more scalable than high-latency models. This is particularly important for businesses that experience spikes in traffic or have a large user base.\n\nIn summary, low latency LLMs are essential for real-time applications, providing a better user experience, enabling quick decision-making, and improving scalability. As the demand for real-time NLP applications continues to grow, the importance of low latency LLMs will only become more critical."
      },
      "finish_reason": "stop",
      "logprobs": null
    }
  ],
  "usage": {
    "prompt_tokens": 24,
    "completion_tokens": 377,
    "total_tokens": 401,
    "prompt_time": 0.009,
    "completion_time": 0.774,
    "total_time": 0.783
  }
}

The groq API documentation was crucial for fetching token usage data. However, I encountered a challenge: the documentation only covered usage stats for regular API calls, not for streaming objects, which Harshil's code utilized. After some trial and error, I discovered that the last stream object contains the usage information and was able to implement the feature.

Implementing the feature

With a clear understanding of the codebase and the API, I proceeded to implement the token usage reporting feature. This involved:

  • Adding the Command-Line Flag to trigger the token usage report.

  • Integrating Token Usage Tracking by modifying the CLI tool to output token usage data to stderr alongside the regular output that the normally tool produces. This was as simple as adding an if condition and then printing coloured information to stderr.

      # Print colored token usage info
  # Ref Doc: https://codehs.com/tutorial/andy/ansi-colors
  if token_usage:
      usage = chunk.x_groq.usage

      formatted_usage = (
          "\n\033[92m"
          "Token Usage:\n"
          "-------------\n"
          f"- Completion Tokens: {usage.completion_tokens}\n"
          f"- Prompt Tokens: {usage.prompt_tokens}\n"
          f"- Total Tokens: {usage.total_tokens}\n\n"
          "Timing:\n"
          "-------\n"
          f"- Completion Time: {usage.completion_time:.3f} seconds\n"
          f"- Prompt Time: {usage.prompt_time:.3f} seconds\n"
          f"- Queue Time: {usage.queue_time:.3f} seconds\n"
          f"- Total Time: {usage.total_time:.3f} seconds\n"
          "\033[0m"
      )

      print(formatted_usage, file=sys.stderr)

  • After implementing the feature, I thoroughly tested it to ensure accuracy and functionality. The new flag was now correctly reporting token usage statistics.

  • Following successful testing, I updated the project documentation to include information about the new feature and instructions on how to use it.

  • I also setup a basic code formatter in the code so that the code is easier to maintain in the future.

More improvements

Another challenge was the spinner in the existing code. The spinner was reloading continuously every time content was printed to the console. I fixed this by changing the way content was printed. Instead of printing chunks of content immediately, I accumulated the content in a variable and printed it all at once, which resolved the issue with the spinner. Here’s the relevant code change:

Before:

if chunk_content:
    if output is None:
        print(chunk_content, end="")
    else:
        content += chunk_content

After:

if chunk_content:
    content += chunk_content

print(content)

This change prevents the spinner from reloading with each print statement and ensures smoother operation.

Submitting a PR

All of these changes were made on my own branch in a forked copy of Harshil’s repository so I had to open a pull request as you usually do not have write access to open-source repositories unless you are one of the admins or maintainers. You can check out my pull request that I opened up once I fixed the issue. I then linked this PR to the issue that opened before.

Harshil liked my work on his code and merged the PR right away. After the PR was merged, I closed the issue that I had opened earlier.

Approving pull requests

As this is a part of the Open Source course that I mentioned before, Harshil made the same contribution to my GitHub repository, github-echo, a command-line tool designed to provide actionable insights about GitHub repositories. He also implemented the token usage reporting feature, which enhances the tool's functionality by tracking and reporting token usage effectively.

You can review the details of the feature request and implementation through the following links:

The process of integrating Harshil's changes was seamless, thanks to our strong collaborative relationship. I tested the updates by fetching his pull request into a separate branch and thoroughly evaluating the new feature. Everything worked flawlessly, so I left a positive comment on his PR, acknowledging the excellent work, and closed the request. I can imagine this process being harder in larger projects and major code changes and I cannot wait to experience that.

Reflection and Learnings

Although the process of contributing to this project wasn’t entirely new to me—thanks to my co-op experience and hackathons—it was my first time collaborating with classmates on such a project. This experience was enjoyable and provided a valuable learning opportunity. I’m eager to tackle more complex projects in the future, especially those with frequently updated main sources of truth. Dealing with more frequent updates and merge conflicts will offer a deeper understanding of git and GitHub, which I look forward to exploring.

That’s all for now! Until next time 😊

osd600Open Source